/* Copyright 2026 Marimo. All rights reserved. */

import { atom } from "jotai";
import { atomWithStorage } from "jotai/utils";
import { isPlatformWindows } from "@/core/hotkeys/shortcuts";
import { jotaiJsonStorage } from "@/utils/storage/jotai";
import { capitalize } from "@/utils/strings";
import type { TypedString } from "@/utils/typed";
import { generateUUID } from "@/utils/uuid";
import type { ExternalAgentSessionId, SessionSupportType } from "./types";

// Types
export type TabId = TypedString<"TabId">;
export type ExternalAgentId =
  | "claude"
  | "gemini"
  | "codex"
  | "opencode"
  | "cursor";

// No agents support loading sessions, so we limit to 1, otherwise
// this is confusing to the user when switching between sessions
const MAX_SESSIONS = 1;

export interface AgentSession {
  tabId: TabId;
  agentId: ExternalAgentId;
  title: string;
  createdAt: number;
  lastUsedAt: number;
  // Store the actual agent session ID for resumption
  externalAgentSessionId: ExternalAgentSessionId | null;
  // Selected model for this session (null = agent default)
  selectedModel: string | null;
}

export interface AgentSessionState {
  sessions: AgentSession[];
  activeTabId: TabId | null;
}

// Constants
const STORAGE_KEY = "marimo:acp:sessions:v1";

// Atoms
export const agentSessionStateAtom = atomWithStorage<AgentSessionState>(
  STORAGE_KEY,
  {
    sessions: [],
    activeTabId: null,
  },
  jotaiJsonStorage,
);

export const selectedTabAtom = atom(
  (get) => {
    const state = get(agentSessionStateAtom);
    if (!state.activeTabId) {
      return null;
    }
    return (
      state.sessions.find((session) => session.tabId === state.activeTabId) ||
      null
    );
  },
  (_get, set, activeTabId: TabId | null) => {
    set(agentSessionStateAtom, (prev) => ({
      ...prev,
      activeTabId: activeTabId,
    }));
  },
);

// Utilities
function generateTabId(): TabId {
  // Our tab ID for internal session management
  return `tab_${generateUUID()}` as TabId;
}

export function truncateTitle(title: string, maxLength = 20): string {
  if (title.length <= maxLength) {
    return title;
  }
  return `${title.slice(0, maxLength - 3)}...`;
}

export function addSession(
  state: AgentSessionState,
  session: {
    agentId: ExternalAgentId;
    firstMessage?: string;
    model?: string | null;
  },
): AgentSessionState {
  const sessionSupport = getAgentSessionSupport(session.agentId);

  const now = Date.now();
  const title = session.firstMessage
    ? truncateTitle(session.firstMessage.trim())
    : `New ${session.agentId} session`;
  const tabId = generateTabId();

  if (sessionSupport === "single") {
    // For single session agents, replace any existing session for this agent
    const existingSessions = state.sessions.filter(
      (s) => s.agentId === session.agentId,
    );
    const otherSessions = state.sessions.filter(
      (s) => s.agentId !== session.agentId,
    );

    if (existingSessions.length > 0) {
      // Replace the existing session (overwrite it)
      const existingSession = existingSessions[0];
      const updatedSession: AgentSession = {
        agentId: session.agentId,
        title,
        createdAt: now,
        lastUsedAt: now,
        tabId: existingSession.tabId, // Keep the same ID to maintain tab reference
        externalAgentSessionId: null, // Clear the external session ID
        selectedModel: session.model ?? existingSession.selectedModel ?? null,
      };

      return {
        ...state,
        sessions: [...otherSessions.slice(0, MAX_SESSIONS - 1), updatedSession],
        activeTabId: updatedSession.tabId,
      };
    }
  }

  // For multiple session agents or when no existing session exists
  return {
    ...state,
    sessions: [
      ...state.sessions.slice(0, MAX_SESSIONS - 1),
      {
        agentId: session.agentId,
        tabId,
        title,
        createdAt: now,
        lastUsedAt: now,
        externalAgentSessionId: null,
        selectedModel: session.model ?? null,
      },
    ],
    activeTabId: tabId,
  };
}

export function removeSession(
  state: AgentSessionState,
  sessionId: TabId,
): AgentSessionState {
  const filteredSessions = state.sessions.filter((s) => s.tabId !== sessionId);
  const newActiveTabId =
    state.activeTabId === sessionId
      ? filteredSessions.length > 0
        ? filteredSessions[filteredSessions.length - 1].tabId
        : null
      : state.activeTabId;

  return {
    sessions: filteredSessions,
    activeTabId: newActiveTabId,
  };
}

export function updateSessionTitle(
  state: AgentSessionState,
  title: string,
): AgentSessionState {
  const selectedTab = state.activeTabId;
  if (!selectedTab) {
    return state;
  }
  return {
    ...state,
    sessions: state.sessions.map((session) =>
      session.tabId === selectedTab
        ? { ...session, title: truncateTitle(title) }
        : session,
    ),
  };
}

export function updateSessionLastUsed(
  state: AgentSessionState,
  sessionId: TabId,
): AgentSessionState {
  return {
    ...state,
    sessions: state.sessions.map((session) =>
      session.tabId === sessionId
        ? { ...session, lastUsedAt: Date.now() }
        : session,
    ),
  };
}

/**
 * Update the sessionId for the currently selected tab
 */
export function updateSessionExternalAgentSessionId(
  state: AgentSessionState,
  externalAgentSessionId: ExternalAgentSessionId,
): AgentSessionState {
  const selectedTab = state.activeTabId;
  if (!selectedTab) {
    return state;
  }
  return {
    ...state,
    sessions: state.sessions.map((session) =>
      session.tabId === selectedTab
        ? { ...session, externalAgentSessionId, lastUsedAt: Date.now() }
        : session,
    ),
  };
}

export function getSessionsByAgent(
  sessions: AgentSession[],
  agentId: ExternalAgentId,
): AgentSession[] {
  return sessions
    .filter((session) => session.agentId === agentId)
    .toSorted((a, b) => b.lastUsedAt - a.lastUsedAt);
}

export function getAllAgentIds(): ExternalAgentId[] {
  return ["claude", "gemini", "codex", "opencode", "cursor"];
}

export function getAgentDisplayName(agentId: ExternalAgentId): string {
  return capitalize(agentId);
}

export function getAgentWebSocketUrl(agentId: ExternalAgentId): string {
  const port = AGENT_CONFIG[agentId].port;
  // Use the current page's hostname so the agent is reachable when
  // marimo is accessed remotely (e.g. via direct IP or reverse proxy).
  const hostname = window.location.hostname;
  const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
  return `${protocol}//${hostname}:${port}/message` as const;
}

interface AgentConfig {
  port: number;
  command: string;
  sessionSupport: SessionSupportType;
  /** One-time setup command the user must run before starting the agent. */
  loginHint?: string;
}

const AGENT_CONFIG: Record<ExternalAgentId, AgentConfig> = {
  claude: {
    port: 3017,
    command: "npx @zed-industries/claude-code-acp",
    sessionSupport: "single",
  },
  gemini: {
    port: 3019,
    command: "npx @google/gemini-cli --experimental-acp",
    sessionSupport: "single",
  },
  codex: {
    port: 3021,
    command: "npx @zed-industries/codex-acp",
    sessionSupport: "single",
  },
  opencode: {
    port: 3023,
    command: "npx opencode-ai acp",
    sessionSupport: "single",
  },
  cursor: {
    port: 3025,
    command: "agent acp",
    sessionSupport: "single",
  },
};

export function getAgentSessionSupport(
  agentId: ExternalAgentId,
): SessionSupportType {
  return AGENT_CONFIG[agentId].sessionSupport;
}

export function getAgentConnectionCommand(agentId: ExternalAgentId): string {
  const port = AGENT_CONFIG[agentId].port;
  const command = AGENT_CONFIG[agentId].command;
  const wrappedCommand = isPlatformWindows() ? `cmd /c ${command}` : command;
  return `npx stdio-to-ws "${wrappedCommand}" --port ${port}`;
}

/**
 * Update the selected model for the currently active session
 */
export function updateSessionModel(
  state: AgentSessionState,
  model: string | null,
): AgentSessionState {
  const selectedTab = state.activeTabId;
  if (!selectedTab) {
    return state;
  }
  return {
    ...state,
    sessions: state.sessions.map((session) =>
      session.tabId === selectedTab
        ? { ...session, selectedModel: model }
        : session,
    ),
  };
}