wraith/src/stores/session.store.ts

import { defineStore } from "pinia";
import { ref, computed } from "vue";
import { invoke } from "@tauri-apps/api/core";
import { listen } from "@tauri-apps/api/event";
import { useConnectionStore } from "@/stores/connection.store";
import type { ThemeDefinition } from "@/components/common/ThemePicker.vue";

export interface Session {
  id: string;
  connectionId: number;
  name: string;
  protocol: "ssh" | "rdp" | "local";
  active: boolean;
  username?: string;
  status: "connected" | "disconnected";
  hasActivity: boolean;
}

export interface TerminalDimensions {
  cols: number;
  rows: number;
}

export const useSessionStore = defineStore("session", () => {
  const sessions = ref<Session[]>([]);
  const activeSessionId = ref<string | null>(null);
  const connecting = ref(false);
  const lastError = ref<string | null>(null);

  /** Active terminal theme — applied to all terminal instances. */
  const activeTheme = ref<ThemeDefinition | null>(null);

  /** Per-session terminal dimensions (cols x rows). */
  const terminalDimensions = ref<Record<string, TerminalDimensions>>({});

  const activeSession = computed(() =>
    sessions.value.find((s) => s.id === activeSessionId.value) ?? null,
  );

  const sessionCount = computed(() => sessions.value.length);

  // Listen for backend close/exit events to update session status
  function setupStatusListeners(sessionId: string): void {
    listen(`ssh:close:${sessionId}`, () => markDisconnected(sessionId));
    listen(`ssh:exit:${sessionId}`, () => markDisconnected(sessionId));
  }

  function markDisconnected(sessionId: string): void {
    const session = sessions.value.find((s) => s.id === sessionId);
    if (session) session.status = "disconnected";
  }

  function activateSession(id: string): void {
    activeSessionId.value = id;
    // Clear activity indicator when switching to tab
    const session = sessions.value.find(s => s.id === id);
    if (session) session.hasActivity = false;
  }

  /** Mark a background tab as having new activity. */
  function markActivity(sessionId: string): void {
    if (sessionId === activeSessionId.value) return; // don't flash the active tab
    const session = sessions.value.find(s => s.id === sessionId);
    if (session) session.hasActivity = true;
  }

  /** Reorder sessions by moving a tab from one index to another. */
  function moveSession(fromIndex: number, toIndex: number): void {
    if (fromIndex === toIndex) return;
    if (fromIndex < 0 || toIndex < 0) return;
    if (fromIndex >= sessions.value.length || toIndex >= sessions.value.length) return;
    const [moved] = sessions.value.splice(fromIndex, 1);
    sessions.value.splice(toIndex, 0, moved);
  }

  async function closeSession(id: string): Promise<void> {
    const idx = sessions.value.findIndex((s) => s.id === id);
    if (idx === -1) return;

    const session = sessions.value[idx];

    // Disconnect the backend session using the protocol-appropriate command
    try {
      if (session.protocol === "local") {
        await invoke("disconnect_pty", { sessionId: session.id });
      } else if (session.protocol === "rdp") {
        await invoke("disconnect_rdp", { sessionId: session.id });
      } else {
        await invoke("disconnect_session", { sessionId: session.id });
      }
    } catch (err) {
      console.error("Failed to disconnect session:", err);
    }

    sessions.value.splice(idx, 1);

    if (activeSessionId.value === id) {
      if (sessions.value.length === 0) {
        activeSessionId.value = null;
      } else {
        const nextIdx = Math.min(idx, sessions.value.length - 1);
        activeSessionId.value = sessions.value[nextIdx].id;
      }
    }
  }

  /** Count how many sessions already exist for this connection (for tab name disambiguation). */
  function sessionCountForConnection(connId: number): number {
    return sessions.value.filter((s) => s.connectionId === connId).length;
  }

  /** Generate a disambiguated tab name like "Asgard", "Asgard (2)", "Asgard (3)". */
  function disambiguatedName(baseName: string, connId: number): string {
    const count = sessionCountForConnection(connId);
    return count === 0 ? baseName : `${baseName} (${count + 1})`;
  }

  /**
   * Connect to a server by connection ID.
   * Multiple sessions to the same host are allowed (MobaXTerm-style).
   * Each gets its own tab with a disambiguated name like "Asgard (2)".
   *
   * For Tauri: we must resolve the connection details ourselves and pass
   * hostname/port/username/password directly to connect_ssh, because the
   * Rust side has no knowledge of connection IDs — the vault owns credentials.
   */
  async function connect(connectionId: number): Promise<void> {
    const connectionStore = useConnectionStore();
    const conn = connectionStore.connections.find((c) => c.id === connectionId);
    if (!conn) return;

    connecting.value = true;
    try {
      if (conn.protocol === "ssh") {
        let sessionId: string;
        let resolvedUsername = "";
        let resolvedPassword = "";

        // If connection has a linked credential, decrypt it from the vault
        if (conn.credentialId) {
          try {
            const allCreds = await invoke<{ id: number; name: string; username: string | null; credentialType: string; sshKeyId: number | null }[]>("list_credentials");
            const cred = allCreds.find((c) => c.id === conn.credentialId);

            if (cred) {
              resolvedUsername = cred.username ?? "";

              if (cred.credentialType === "ssh_key" && cred.sshKeyId) {
                // SSH key auth — decrypt key from vault
                const [privateKey, passphrase] = await invoke<[string, string]>("decrypt_ssh_key", { sshKeyId: cred.sshKeyId });
                sessionId = await invoke<string>("connect_ssh_with_key", {
                  hostname: conn.hostname,
                  port: conn.port,
                  username: resolvedUsername,
                  privateKeyPem: privateKey,
                  passphrase: passphrase || null,
                  cols: 120,
                  rows: 40,
                });

                sessions.value.push({
                  id: sessionId,
                  connectionId,
                  name: disambiguatedName(conn.name, connectionId),
                  protocol: "ssh",
                  active: true,
                  username: resolvedUsername,
                  status: "connected",
                  hasActivity: false,
                });
                setupStatusListeners(sessionId);
                activeSessionId.value = sessionId;
                return; // early return — key auth handled
              } else {
                // Password auth — decrypt password from vault
                resolvedPassword = await invoke<string>("decrypt_password", { credentialId: cred.id });
              }
            }
          } catch (credErr) {
            console.warn("Failed to resolve credential, will prompt:", credErr);
          }
        }

        try {
          if (!resolvedUsername) {
            // No credential linked — prompt immediately
            throw new Error("NO_CREDENTIALS");
          }
          sessionId = await invoke<string>("connect_ssh", {
            hostname: conn.hostname,
            port: conn.port,
            username: resolvedUsername,
            password: resolvedPassword,
            cols: 120,
            rows: 40,
          });
        } catch (sshErr: unknown) {
          const errMsg = sshErr instanceof Error
            ? sshErr.message
            : typeof sshErr === "string"
              ? sshErr
              : String(sshErr);

          // If no credentials or auth failed, prompt for username/password
          const errLower = errMsg.toLowerCase();
          if (errLower.includes("no_credentials") || errLower.includes("unable to authenticate") || errLower.includes("authentication") || errLower.includes("rejected")) {
            const username = window.prompt(`Username for ${conn.hostname}:`, resolvedUsername || "root");
            if (!username) throw new Error("Connection cancelled");
            const password = window.prompt(`Password for ${username}@${conn.hostname}:`);
            if (password === null) throw new Error("Connection cancelled");

            resolvedUsername = username;
            sessionId = await invoke<string>("connect_ssh", {
              hostname: conn.hostname,
              port: conn.port,
              username,
              password,
              cols: 120,
              rows: 40,
            });
          } else {
            throw sshErr;
          }
        }

        sessions.value.push({
          id: sessionId,
          connectionId,
          name: disambiguatedName(conn.name, connectionId),
          protocol: "ssh",
          active: true,
          username: resolvedUsername,
          status: "connected",
          hasActivity: false,
        });
        setupStatusListeners(sessionId);
        activeSessionId.value = sessionId;
      } else if (conn.protocol === "rdp") {
        let username = "";
        let password = "";
        let domain = "";

        // Try vault credentials first (same as SSH path)
        if (conn.credentialId) {
          try {
            const allCreds = await invoke<{ id: number; name: string; username: string | null; domain: string | null; credentialType: string; sshKeyId: number | null }[]>("list_credentials");
            const cred = allCreds.find((c) => c.id === conn.credentialId);
            if (cred && cred.credentialType === "password") {
              username = cred.username ?? "";
              domain = cred.domain ?? "";
              password = await invoke<string>("decrypt_password", { credentialId: cred.id });
            }
          } catch (credErr) {
            console.warn("Failed to resolve RDP credential from vault:", credErr);
          }
        }

        // Fall back to connection options JSON if vault didn't provide creds
        if (!username && conn.options) {
          try {
            const opts = JSON.parse(conn.options);
            if (opts?.username) username = opts.username;
            if (opts?.password) password = opts.password;
            if (opts?.domain) domain = opts.domain;
          } catch {
            // ignore malformed options
          }
        }

        let sessionId: string;
        try {
          sessionId = await invoke<string>("connect_rdp", {
            config: {
              hostname: conn.hostname,
              port: conn.port,
              username,
              password,
              domain,
              width: 1920,
              height: 1080,
            },
          });
        } catch (rdpErr: unknown) {
          const errMsg =
            rdpErr instanceof Error
              ? rdpErr.message
              : typeof rdpErr === "string"
                ? rdpErr
                : String(rdpErr);

          // If credentials are missing or rejected, prompt the operator
          if (
            errMsg.includes("NO_CREDENTIALS") ||
            errMsg.includes("authentication") ||
            errMsg.includes("logon failure")
          ) {
            const promptedUsername = prompt(
              `Username for ${conn.hostname}:`,
              "Administrator",
            );
            if (!promptedUsername) throw new Error("Connection cancelled");
            const promptedPassword = prompt(
              `Password for ${promptedUsername}@${conn.hostname}:`,
            );
            if (promptedPassword === null) throw new Error("Connection cancelled");
            const promptedDomain = prompt(`Domain (leave blank if none):`, "") ?? "";

            username = promptedUsername;

            sessionId = await invoke<string>("connect_rdp", {
              config: {
                hostname: conn.hostname,
                port: conn.port,
                username: promptedUsername,
                password: promptedPassword,
                domain: promptedDomain,
                width: 1920,
                height: 1080,
              },
            });
          } else {
            throw rdpErr;
          }
        }

        sessions.value.push({
          id: sessionId,
          connectionId,
          name: disambiguatedName(conn.name, connectionId),
          protocol: "rdp",
          active: true,
          username,
          status: "connected",
          hasActivity: false,
        });
        activeSessionId.value = sessionId;
      }
    } catch (err: unknown) {
      const msg = err instanceof Error ? err.message : typeof err === "string" ? err : String(err);
      console.error("Connection failed:", msg);
      lastError.value = msg;
      // Show error as native alert so it's visible without DevTools
      alert(`Connection failed: ${msg}`);
    } finally {
      connecting.value = false;
    }
  }

  /** Spawn a local shell as a full-size tab. */
  async function spawnLocalTab(shellName: string, shellPath: string): Promise<void> {
    try {
      const sessionId = await invoke<string>("spawn_local_shell", {
        shellPath,
        cols: 120,
        rows: 40,
      });

      sessions.value.push({
        id: sessionId,
        connectionId: 0,
        name: shellName,
        protocol: "local",
        active: true,
        status: "connected",
        hasActivity: false,
      });

      // Listen for PTY close
      listen(`pty:close:${sessionId}`, () => markDisconnected(sessionId));

      activeSessionId.value = sessionId;
    } catch (err: unknown) {
      const msg = err instanceof Error ? err.message : String(err);
      alert(`Failed to spawn local shell: ${msg}`);
    }
  }

  /** Apply a theme to all active terminal instances. */
  function setTheme(theme: ThemeDefinition): void {
    activeTheme.value = theme;
  }

  /** Update the recorded dimensions for a terminal session. */
  function setTerminalDimensions(sessionId: string, cols: number, rows: number): void {
    terminalDimensions.value[sessionId] = { cols, rows };
  }

  /** Get the dimensions for the active session, or null if not tracked yet. */
  const activeDimensions = computed<TerminalDimensions | null>(() => {
    if (!activeSessionId.value) return null;
    return terminalDimensions.value[activeSessionId.value] ?? null;
  });

  return {
    sessions,
    activeSessionId,
    activeSession,
    sessionCount,
    connecting,
    lastError,
    activeTheme,
    terminalDimensions,
    activeDimensions,
    activateSession,
    closeSession,
    connect,
    spawnLocalTab,
    moveSession,
    markActivity,
    setTheme,
    setTerminalDimensions,
  };
});