export type WsRequest =
  | {
      type: "chat";
      messages: Message[];
      config: ProviderConfig;
    }
  | {
      type: "cancel";
    }
  | {
      type: "permission_response";
      request_id: string;
      approved: boolean;
      always_allow: boolean;
    }
  | { type: "ping" }
  | {
      type: "side_question";
      question: string;
      context_messages: Message[];
      config: ProviderConfig;
    };

export interface AgentAssignment {
  agent_name: string;
  model: string | null;
  status: string;
}

export interface PipelineStageItem {
  story_id: string;
  name: string | null;
  error: string | null;
  merge_failure: string | null;
  agent: AgentAssignment | null;
  review_hold: boolean | null;
  qa: string | null;
}

export interface PipelineState {
  backlog: PipelineStageItem[];
  current: PipelineStageItem[];
  qa: PipelineStageItem[];
  merge: PipelineStageItem[];
  done: PipelineStageItem[];
}

export type WsResponse =
  | { type: "token"; content: string }
  | { type: "update"; messages: Message[] }
  | { type: "session_id"; session_id: string }
  | { type: "error"; message: string }
  | {
      type: "pipeline_state";
      backlog: PipelineStageItem[];
      current: PipelineStageItem[];
      qa: PipelineStageItem[];
      merge: PipelineStageItem[];
      done: PipelineStageItem[];
    }
  | {
      type: "permission_request";
      request_id: string;
      tool_name: string;
      tool_input: Record<string, unknown>;
    }
  | { type: "tool_activity"; tool_name: string }
  | {
      type: "reconciliation_progress";
      story_id: string;
      status: string;
      message: string;
    }
  /** `.story_kit/project.toml` was modified; re-fetch the agent roster. */
  | { type: "agent_config_changed" }
  /** An agent started, stopped, or changed state; re-fetch agent list. */
  | { type: "agent_state_changed" }
  | { type: "tool_activity"; tool_name: string }
  /** Heartbeat response confirming the connection is alive. */
  | { type: "pong" }
  /** Sent on connect when the project still needs onboarding (specs are placeholders). */
  | { type: "onboarding_status"; needs_onboarding: boolean }
  /** Streaming thinking token from an extended-thinking block, separate from regular text. */
  | { type: "thinking_token"; content: string }
  /** Streaming token from a /btw side question response. */
  | { type: "side_question_token"; content: string }
  /** Final signal that the /btw side question has been fully answered. */
  | { type: "side_question_done"; response: string }
  /** A single server log entry (bulk on connect, then live). */
  | { type: "log_entry"; timestamp: string; level: string; message: string };

export interface ProviderConfig {
  provider: string;
  model: string;
  base_url?: string;
  enable_tools?: boolean;
  session_id?: string;
}

export type Role = "system" | "user" | "assistant" | "tool";

export interface ToolCall {
  id?: string;
  type: string;
  function: {
    name: string;
    arguments: string;
  };
}

export interface Message {
  role: Role;
  content: string;
  tool_calls?: ToolCall[];
  tool_call_id?: string;
}

export interface WorkItemContent {
  content: string;
  stage: string;
  name: string | null;
  agent: string | null;
}

export interface TestCaseResult {
  name: string;
  status: "pass" | "fail";
  details: string | null;
}

export interface TestResultsResponse {
  unit: TestCaseResult[];
  integration: TestCaseResult[];
}

export interface FileEntry {
  name: string;
  kind: "file" | "dir";
}

export interface SearchResult {
  path: string;
  matches: number;
}

export interface AgentCostEntry {
  agent_name: string;
  model: string | null;
  input_tokens: number;
  output_tokens: number;
  cache_creation_input_tokens: number;
  cache_read_input_tokens: number;
  total_cost_usd: number;
}

export interface TokenCostResponse {
  total_cost_usd: number;
  agents: AgentCostEntry[];
}

export interface TokenUsageRecord {
  story_id: string;
  agent_name: string;
  model: string | null;
  timestamp: string;
  input_tokens: number;
  output_tokens: number;
  cache_creation_input_tokens: number;
  cache_read_input_tokens: number;
  total_cost_usd: number;
}

export interface AllTokenUsageResponse {
  records: TokenUsageRecord[];
}

export interface CommandOutput {
  stdout: string;
  stderr: string;
  exit_code: number;
}

declare const __STORKIT_PORT__: string;

const DEFAULT_API_BASE = "/api";
const DEFAULT_WS_PATH = "/ws";

export function resolveWsHost(
  isDev: boolean,
  envPort: string | undefined,
  locationHost: string,
): string {
  return isDev ? `127.0.0.1:${envPort || "3001"}` : locationHost;
}

function buildApiUrl(path: string, baseUrl = DEFAULT_API_BASE): string {
  return `${baseUrl}${path}`;
}

async function requestJson<T>(
  path: string,
  options: RequestInit = {},
  baseUrl = DEFAULT_API_BASE,
): Promise<T> {
  const res = await fetch(buildApiUrl(path, baseUrl), {
    headers: {
      "Content-Type": "application/json",
      ...(options.headers ?? {}),
    },
    ...options,
  });

  if (!res.ok) {
    const text = await res.text();
    throw new Error(text || `Request failed (${res.status})`);
  }

  return res.json() as Promise<T>;
}

export const api = {
  getCurrentProject(baseUrl?: string) {
    return requestJson<string | null>("/project", {}, baseUrl);
  },
  getKnownProjects(baseUrl?: string) {
    return requestJson<string[]>("/projects", {}, baseUrl);
  },
  forgetKnownProject(path: string, baseUrl?: string) {
    return requestJson<boolean>(
      "/projects/forget",
      { method: "POST", body: JSON.stringify({ path }) },
      baseUrl,
    );
  },
  openProject(path: string, baseUrl?: string) {
    return requestJson<string>(
      "/project",
      { method: "POST", body: JSON.stringify({ path }) },
      baseUrl,
    );
  },
  closeProject(baseUrl?: string) {
    return requestJson<boolean>("/project", { method: "DELETE" }, baseUrl);
  },
  getModelPreference(baseUrl?: string) {
    return requestJson<string | null>("/model", {}, baseUrl);
  },
  setModelPreference(model: string, baseUrl?: string) {
    return requestJson<boolean>(
      "/model",
      { method: "POST", body: JSON.stringify({ model }) },
      baseUrl,
    );
  },
  getOllamaModels(baseUrlParam?: string, baseUrl?: string) {
    const url = new URL(
      buildApiUrl("/ollama/models", baseUrl),
      window.location.origin,
    );
    if (baseUrlParam) {
      url.searchParams.set("base_url", baseUrlParam);
    }
    return requestJson<string[]>(url.pathname + url.search, {}, "");
  },
  getAnthropicApiKeyExists(baseUrl?: string) {
    return requestJson<boolean>("/anthropic/key/exists", {}, baseUrl);
  },
  getAnthropicModels(baseUrl?: string) {
    return requestJson<string[]>("/anthropic/models", {}, baseUrl);
  },
  setAnthropicApiKey(api_key: string, baseUrl?: string) {
    return requestJson<boolean>(
      "/anthropic/key",
      { method: "POST", body: JSON.stringify({ api_key }) },
      baseUrl,
    );
  },
  readFile(path: string, baseUrl?: string) {
    return requestJson<string>(
      "/fs/read",
      { method: "POST", body: JSON.stringify({ path }) },
      baseUrl,
    );
  },
  writeFile(path: string, content: string, baseUrl?: string) {
    return requestJson<boolean>(
      "/fs/write",
      { method: "POST", body: JSON.stringify({ path, content }) },
      baseUrl,
    );
  },
  listDirectory(path: string, baseUrl?: string) {
    return requestJson<FileEntry[]>(
      "/fs/list",
      { method: "POST", body: JSON.stringify({ path }) },
      baseUrl,
    );
  },
  listDirectoryAbsolute(path: string, baseUrl?: string) {
    return requestJson<FileEntry[]>(
      "/io/fs/list/absolute",
      { method: "POST", body: JSON.stringify({ path }) },
      baseUrl,
    );
  },
  createDirectoryAbsolute(path: string, baseUrl?: string) {
    return requestJson<boolean>(
      "/io/fs/create/absolute",
      { method: "POST", body: JSON.stringify({ path }) },
      baseUrl,
    );
  },
  getHomeDirectory(baseUrl?: string) {
    return requestJson<string>("/io/fs/home", {}, baseUrl);
  },
  listProjectFiles(baseUrl?: string) {
    return requestJson<string[]>("/io/fs/files", {}, baseUrl);
  },
  searchFiles(query: string, baseUrl?: string) {
    return requestJson<SearchResult[]>(
      "/fs/search",
      { method: "POST", body: JSON.stringify({ query }) },
      baseUrl,
    );
  },
  execShell(command: string, args: string[], baseUrl?: string) {
    return requestJson<CommandOutput>(
      "/shell/exec",
      { method: "POST", body: JSON.stringify({ command, args }) },
      baseUrl,
    );
  },
  cancelChat(baseUrl?: string) {
    return requestJson<boolean>("/chat/cancel", { method: "POST" }, baseUrl);
  },
  getWorkItemContent(storyId: string, baseUrl?: string) {
    return requestJson<WorkItemContent>(
      `/work-items/${encodeURIComponent(storyId)}`,
      {},
      baseUrl,
    );
  },
  getTestResults(storyId: string, baseUrl?: string) {
    return requestJson<TestResultsResponse | null>(
      `/work-items/${encodeURIComponent(storyId)}/test-results`,
      {},
      baseUrl,
    );
  },
  getTokenCost(storyId: string, baseUrl?: string) {
    return requestJson<TokenCostResponse>(
      `/work-items/${encodeURIComponent(storyId)}/token-cost`,
      {},
      baseUrl,
    );
  },
  getAllTokenUsage(baseUrl?: string) {
    return requestJson<AllTokenUsageResponse>("/token-usage", {}, baseUrl);
  },
  /** Trigger a server rebuild and restart. */
  rebuildAndRestart() {
    return callMcpTool("rebuild_and_restart", {});
  },
  /** Approve a story in QA, moving it to merge. */
  approveQa(storyId: string) {
    return callMcpTool("approve_qa", { story_id: storyId });
  },
  /** Reject a story in QA, moving it back to current with notes. */
  rejectQa(storyId: string, notes: string) {
    return callMcpTool("reject_qa", { story_id: storyId, notes });
  },
  /** Launch the QA app for a story's worktree. */
  launchQaApp(storyId: string) {
    return callMcpTool("launch_qa_app", { story_id: storyId });
  },
  /** Delete a story from the pipeline, stopping any running agent and removing the worktree. */
  deleteStory(storyId: string) {
    return callMcpTool("delete_story", { story_id: storyId });
  },
};

async function callMcpTool(
  toolName: string,
  args: Record<string, unknown>,
): Promise<string> {
  const res = await fetch("/mcp", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({
      jsonrpc: "2.0",
      id: 1,
      method: "tools/call",
      params: { name: toolName, arguments: args },
    }),
  });
  const json = await res.json();
  if (json.error) {
    throw new Error(json.error.message);
  }
  const text = json.result?.content?.[0]?.text ?? "";
  return text;
}

export class ChatWebSocket {
  private static sharedSocket: WebSocket | null = null;
  private static refCount = 0;
  private socket?: WebSocket;
  private onToken?: (content: string) => void;
  private onThinkingToken?: (content: string) => void;
  private onUpdate?: (messages: Message[]) => void;
  private onSessionId?: (sessionId: string) => void;
  private onError?: (message: string) => void;
  private onPipelineState?: (state: PipelineState) => void;
  private onPermissionRequest?: (
    requestId: string,
    toolName: string,
    toolInput: Record<string, unknown>,
  ) => void;
  private onActivity?: (toolName: string) => void;
  private onReconciliationProgress?: (
    storyId: string,
    status: string,
    message: string,
  ) => void;
  private onAgentConfigChanged?: () => void;
  private onAgentStateChanged?: () => void;
  private onOnboardingStatus?: (needsOnboarding: boolean) => void;
  private onSideQuestionToken?: (content: string) => void;
  private onSideQuestionDone?: (response: string) => void;
  private onLogEntry?: (
    timestamp: string,
    level: string,
    message: string,
  ) => void;
  private onConnected?: () => void;
  private connected = false;
  private closeTimer?: number;
  private wsPath = DEFAULT_WS_PATH;
  private reconnectTimer?: number;
  private reconnectDelay = 1000;
  private shouldReconnect = false;
  private heartbeatInterval?: number;
  private heartbeatTimeout?: number;
  private static readonly HEARTBEAT_INTERVAL = 30_000;
  private static readonly HEARTBEAT_TIMEOUT = 5_000;

  private _startHeartbeat(): void {
    this._stopHeartbeat();
    this.heartbeatInterval = window.setInterval(() => {
      if (!this.socket || this.socket.readyState !== WebSocket.OPEN) return;
      const ping: WsRequest = { type: "ping" };
      this.socket.send(JSON.stringify(ping));
      this.heartbeatTimeout = window.setTimeout(() => {
        // No pong received within timeout; close socket to trigger reconnect.
        this.socket?.close();
      }, ChatWebSocket.HEARTBEAT_TIMEOUT);
    }, ChatWebSocket.HEARTBEAT_INTERVAL);
  }

  private _stopHeartbeat(): void {
    window.clearInterval(this.heartbeatInterval);
    window.clearTimeout(this.heartbeatTimeout);
    this.heartbeatInterval = undefined;
    this.heartbeatTimeout = undefined;
  }

  private _buildWsUrl(): string {
    const protocol = window.location.protocol === "https:" ? "wss" : "ws";
    const wsHost = resolveWsHost(
      import.meta.env.DEV,
      typeof __STORKIT_PORT__ !== "undefined" ? __STORKIT_PORT__ : undefined,
      window.location.host,
    );
    return `${protocol}://${wsHost}${this.wsPath}`;
  }

  private _attachHandlers(): void {
    if (!this.socket) return;
    this.socket.onopen = () => {
      this.reconnectDelay = 1000;
      this._startHeartbeat();
      this.onConnected?.();
    };
    this.socket.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data) as WsResponse;
        if (data.type === "token") this.onToken?.(data.content);
        if (data.type === "thinking_token")
          this.onThinkingToken?.(data.content);
        if (data.type === "update") this.onUpdate?.(data.messages);
        if (data.type === "session_id") this.onSessionId?.(data.session_id);
        if (data.type === "error") this.onError?.(data.message);
        if (data.type === "pipeline_state")
          this.onPipelineState?.({
            backlog: data.backlog,
            current: data.current,
            qa: data.qa,
            merge: data.merge,
            done: data.done,
          });
        if (data.type === "permission_request")
          this.onPermissionRequest?.(
            data.request_id,
            data.tool_name,
            data.tool_input,
          );
        if (data.type === "tool_activity") this.onActivity?.(data.tool_name);
        if (data.type === "reconciliation_progress")
          this.onReconciliationProgress?.(
            data.story_id,
            data.status,
            data.message,
          );
        if (data.type === "agent_config_changed") this.onAgentConfigChanged?.();
        if (data.type === "agent_state_changed") this.onAgentStateChanged?.();
        if (data.type === "onboarding_status")
          this.onOnboardingStatus?.(data.needs_onboarding);
        if (data.type === "side_question_token")
          this.onSideQuestionToken?.(data.content);
        if (data.type === "side_question_done")
          this.onSideQuestionDone?.(data.response);
        if (data.type === "log_entry")
          this.onLogEntry?.(data.timestamp, data.level, data.message);
        if (data.type === "pong") {
          window.clearTimeout(this.heartbeatTimeout);
          this.heartbeatTimeout = undefined;
        }
      } catch (err) {
        this.onError?.(String(err));
      }
    };
    this.socket.onerror = () => {
      this.onError?.("WebSocket error");
    };
    this.socket.onclose = () => {
      if (this.shouldReconnect && this.connected) {
        this._scheduleReconnect();
      }
    };
  }

  private _scheduleReconnect(): void {
    window.clearTimeout(this.reconnectTimer);
    const delay = this.reconnectDelay;
    this.reconnectDelay = Math.min(this.reconnectDelay * 2, 30000);
    this.reconnectTimer = window.setTimeout(() => {
      this.reconnectTimer = undefined;
      const wsUrl = this._buildWsUrl();
      ChatWebSocket.sharedSocket = new WebSocket(wsUrl);
      this.socket = ChatWebSocket.sharedSocket;
      this._attachHandlers();
    }, delay);
  }

  connect(
    handlers: {
      onToken?: (content: string) => void;
      onThinkingToken?: (content: string) => void;
      onUpdate?: (messages: Message[]) => void;
      onSessionId?: (sessionId: string) => void;
      onError?: (message: string) => void;
      onPipelineState?: (state: PipelineState) => void;
      onPermissionRequest?: (
        requestId: string,
        toolName: string,
        toolInput: Record<string, unknown>,
      ) => void;
      onActivity?: (toolName: string) => void;
      onReconciliationProgress?: (
        storyId: string,
        status: string,
        message: string,
      ) => void;
      onAgentConfigChanged?: () => void;
      onAgentStateChanged?: () => void;
      onOnboardingStatus?: (needsOnboarding: boolean) => void;
      onSideQuestionToken?: (content: string) => void;
      onSideQuestionDone?: (response: string) => void;
      onLogEntry?: (timestamp: string, level: string, message: string) => void;
      onConnected?: () => void;
    },
    wsPath = DEFAULT_WS_PATH,
  ) {
    this.onToken = handlers.onToken;
    this.onThinkingToken = handlers.onThinkingToken;
    this.onUpdate = handlers.onUpdate;
    this.onSessionId = handlers.onSessionId;
    this.onError = handlers.onError;
    this.onPipelineState = handlers.onPipelineState;
    this.onPermissionRequest = handlers.onPermissionRequest;
    this.onActivity = handlers.onActivity;
    this.onReconciliationProgress = handlers.onReconciliationProgress;
    this.onAgentConfigChanged = handlers.onAgentConfigChanged;
    this.onAgentStateChanged = handlers.onAgentStateChanged;
    this.onOnboardingStatus = handlers.onOnboardingStatus;
    this.onSideQuestionToken = handlers.onSideQuestionToken;
    this.onSideQuestionDone = handlers.onSideQuestionDone;
    this.onLogEntry = handlers.onLogEntry;
    this.onConnected = handlers.onConnected;
    this.wsPath = wsPath;
    this.shouldReconnect = true;

    if (this.connected) {
      return;
    }
    this.connected = true;
    ChatWebSocket.refCount += 1;

    if (
      !ChatWebSocket.sharedSocket ||
      ChatWebSocket.sharedSocket.readyState === WebSocket.CLOSED ||
      ChatWebSocket.sharedSocket.readyState === WebSocket.CLOSING
    ) {
      const wsUrl = this._buildWsUrl();
      ChatWebSocket.sharedSocket = new WebSocket(wsUrl);
    }
    this.socket = ChatWebSocket.sharedSocket;
    this._attachHandlers();
  }

  sendChat(messages: Message[], config: ProviderConfig) {
    this.send({ type: "chat", messages, config });
  }

  sendSideQuestion(
    question: string,
    contextMessages: Message[],
    config: ProviderConfig,
  ) {
    this.send({
      type: "side_question",
      question,
      context_messages: contextMessages,
      config,
    });
  }

  cancel() {
    this.send({ type: "cancel" });
  }

  sendPermissionResponse(
    requestId: string,
    approved: boolean,
    alwaysAllow = false,
  ) {
    this.send({
      type: "permission_response",
      request_id: requestId,
      approved,
      always_allow: alwaysAllow,
    });
  }

  close() {
    this.shouldReconnect = false;
    this._stopHeartbeat();
    window.clearTimeout(this.reconnectTimer);
    this.reconnectTimer = undefined;

    if (!this.connected) return;
    this.connected = false;
    ChatWebSocket.refCount = Math.max(0, ChatWebSocket.refCount - 1);

    if (import.meta.env.DEV) {
      if (this.closeTimer) {
        window.clearTimeout(this.closeTimer);
      }
      this.closeTimer = window.setTimeout(() => {
        if (ChatWebSocket.refCount === 0) {
          ChatWebSocket.sharedSocket?.close();
          ChatWebSocket.sharedSocket = null;
        }
        this.socket = ChatWebSocket.sharedSocket ?? undefined;
        this.closeTimer = undefined;
      }, 250);
      return;
    }

    if (ChatWebSocket.refCount === 0) {
      ChatWebSocket.sharedSocket?.close();
      ChatWebSocket.sharedSocket = null;
    }
    this.socket = ChatWebSocket.sharedSocket ?? undefined;
  }

  private send(payload: WsRequest) {
    if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
      this.onError?.("WebSocket is not connected");
      return;
    }
    this.socket.send(JSON.stringify(payload));
  }
}