import { useReducer } from "react";
import type { ServerMessage } from "@agent-client/shared";

export type ToolEvent = {
  kind: "tool";
  tool_use_id: string;
  tool: string;
  input: unknown;
  output?: string;
  is_error?: boolean;
  pending: boolean;
};

export type TextEvent = {
  kind: "text";
  text: string;
};

export type TurnEvent = TextEvent | ToolEvent;

export type Turn = {
  id: string;
  prompt: string;
  events: TurnEvent[];
  status: "running" | "done" | "error";
  summary?: string;
  error?: string;
  sessionId?: string;
  cost?: number;
};

export type ChatState = {
  turns: Turn[];
};

type Action =
  | { kind: "user_send"; id: string; prompt: string }
  | { kind: "server"; msg: ServerMessage }
  | { kind: "reset" };

function reducer(state: ChatState, action: Action): ChatState {
  switch (action.kind) {
    case "reset":
      return { turns: [] };

    case "user_send":
      return {
        turns: [
          ...state.turns,
          {
            id: action.id,
            prompt: action.prompt,
            events: [],
            status: "running",
          },
        ],
      };

    case "server": {
      const msg = action.msg;
      // ready / errors w/o id are global — drop in v0.1 (console-visible via the
      // ws client, not surfaced in chat history).
      if (msg.type === "ready") return state;
      if (!("id" in msg) || !msg.id) return state;

      const idx = state.turns.findIndex((t) => t.id === msg.id);
      if (idx < 0) return state;

      const turn = state.turns[idx];
      const next = updateTurn(turn, msg);
      if (next === turn) return state;
      return { turns: state.turns.with(idx, next) };
    }
  }
}

// Returns a new Turn (with new events ref only if events actually change), or
// the same reference if the message is a no-op for this turn. Stable refs keep
// downstream React.memo happy and avoid re-parsing markdown on every token.
function updateTurn(turn: Turn, msg: Extract<ServerMessage, { id?: string }>): Turn {
  switch (msg.type) {
    case "turn_started":
      return turn.sessionId === msg.session_id
        ? turn
        : { ...turn, sessionId: msg.session_id };

    case "assistant_text": {
      const last = turn.events[turn.events.length - 1];
      const events =
        last && last.kind === "text"
          ? turn.events.with(turn.events.length - 1, {
              kind: "text" as const,
              text: last.text + msg.text,
            })
          : [...turn.events, { kind: "text" as const, text: msg.text }];
      return { ...turn, events };
    }

    case "tool_use":
      return {
        ...turn,
        events: [
          ...turn.events,
          {
            kind: "tool" as const,
            tool_use_id: msg.tool_use_id,
            tool: msg.tool,
            input: msg.input,
            pending: true,
          },
        ],
      };

    case "tool_result": {
      const i = turn.events.findIndex(
        (e) => e.kind === "tool" && e.tool_use_id === msg.tool_use_id,
      );
      if (i < 0) return turn;
      const t = turn.events[i] as ToolEvent;
      return {
        ...turn,
        events: turn.events.with(i, {
          ...t,
          output: msg.output,
          is_error: msg.is_error,
          pending: false,
        }),
      };
    }

    case "turn_done":
      return {
        ...turn,
        status: msg.success ? "done" : "error",
        summary: msg.summary,
        cost: msg.usage?.total_cost_usd,
      };

    case "error":
      return { ...turn, status: "error", error: msg.message };

    default:
      return turn;
  }
}

export function useChatStore() {
  return useReducer(reducer, { turns: [] });
}