"use client";

import { useState, useRef, useCallback, useEffect } from "react";
import { streamAgui } from "@/lib/agui-client";
import { normalizeAssistantMarkdownForRender } from "@/lib/chat-markdown";
import { publicConfig } from "@/lib/public-config";
import type {
  AguiEvent,
  AguiMessage,
  AguiToolCall,
  RunAgentInput,
} from "@optage/sdk/client";
import {
  isAgentTool,
  type ChatMessage,
  type InterruptPayload,
  type AnalysisStateSnapshot,
  type TextBlock,
  type ToolCallBlock,
} from "@optage/sdk/client";

export type { InterruptPayload } from "@optage/sdk/client";

export interface UseAgentOptions {
  onMessagesChange?: (
    messages: ChatMessage[],
    history: AguiMessage[],
    analysisState?: AnalysisStateSnapshot | null,
    pendingInterrupt?: InterruptPayload | null,
  ) => void;
  repositoryId?: string | null;
}

export interface UseAgentReturn {
  messages: ChatMessage[];
  isStreaming: boolean;
  error: string | null;
  pendingInterrupt: InterruptPayload | null;
  analysisState: AnalysisStateSnapshot | null;
  threadId: string;
  sendMessage: (text: string) => Promise<void>;
  resolveInterrupt: (approved: boolean, reason?: string) => Promise<void>;
  resetSession: (opts: {
    messages: ChatMessage[];
    history: AguiMessage[];
    threadId: string;
    analysisState?: AnalysisStateSnapshot | null;
    pendingInterrupt?: InterruptPayload | null;
  }) => void;
}

const AGENT_URL = "/api/agent";
const COMPACT_URL = "/api/agent/compact";

function normalizeDisplayedAgentError(message: string): string {
  if (/claude code process aborted by user/i.test(message)) {
    return "Agent run was cancelled before completion.";
  }

  if (/aborted.*timeout|timed out while gathering context|server time limit/i.test(message)) {
    return "Agent timed out while gathering context. The run exceeded the server time limit.";
  }

  return message;
}

export function useAgent(options: UseAgentOptions = {}): UseAgentReturn {
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [isStreaming, setIsStreaming] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [pendingInterrupt, setPendingInterrupt] =
    useState<InterruptPayload | null>(null);
  const [analysisState, setAnalysisState] =
    useState<AnalysisStateSnapshot | null>(null);
  const historyRef = useRef<AguiMessage[]>([]);
  const messagesRef = useRef<ChatMessage[]>([]);
  const threadIdRef = useRef<string>(crypto.randomUUID());
  const abortRef = useRef<AbortController | null>(null);

  const addMessage = useCallback((msg: ChatMessage) => {
    setMessages((prev) => [...(prev || []), msg]);
  }, []);

  const updateMessage = useCallback(
    (msgId: string, updater: (m: ChatMessage) => ChatMessage) => {
      setMessages((prev) =>
        (prev || []).map((m) => (m.id === msgId ? updater(m) : m)),
      );
    },
    [],
  );

  const addBlockToMessage = useCallback(
    (msgId: string, block: ChatMessage["blocks"][number]) => {
      updateMessage(msgId, (m) => ({
        ...m,
        blocks: [...m.blocks, block],
      }));
    },
    [updateMessage],
  );

  const updateBlockInMessage = useCallback(
    <T extends ChatMessage["blocks"][number]>(
      msgId: string,
      blockId: string,
      updater: (b: T) => T,
    ) => {
      updateMessage(msgId, (m) => ({
        ...m,
        blocks: m.blocks.map((b) =>
          b.id === blockId ? updater(b as T) : b,
        ),
      }));
    },
    [updateMessage],
  );

  const processStream = useCallback(
    async (input: RunAgentInput) => {
      setIsStreaming(true);
      setError(null);

      const controller = new AbortController();
      abortRef.current = controller;

      let assistantMsgId: string | null = null;
      let currentTextBlockId: string | null = null;
      let completedTextBlockId: string | null = null;
      let currentText = "";

      const toolCallMap = new Map<
        string,
        { blockId: string; argsJson: string; name: string; parentToolCallId?: string }
      >();

      let activeAgentToolCallId: string | null = null;
      let activeAgentBlockId: string | null = null;

      const updateToolBlock = (
        tc: { blockId: string; parentToolCallId?: string },
        updater: (b: ToolCallBlock) => ToolCallBlock,
      ) => {
        if (!assistantMsgId) return;
        if (tc.parentToolCallId && activeAgentBlockId) {
          updateBlockInMessage<ToolCallBlock>(
            assistantMsgId,
            activeAgentBlockId,
            (parent) => ({
              ...parent,
              subAgentToolCalls: (parent.subAgentToolCalls || []).map((sub) =>
                sub.id === tc.blockId ? updater(sub) : sub,
              ),
            }),
          );
        } else {
          updateBlockInMessage<ToolCallBlock>(assistantMsgId, tc.blockId, updater);
        }
      };

      const toolCallsForHistory: AguiToolCall[] = [];
      const toolResultsForHistory: AguiMessage[] = [];
      let lastMessageId = "";

      const ensureAssistantMsg = (): string => {
        if (!assistantMsgId) {
          assistantMsgId = `asst_${crypto.randomUUID()}`;
          addMessage({
            id: assistantMsgId,
            role: "assistant",
            blocks: [],
            timestamp: Date.now(),
          });
        }
        return assistantMsgId;
      };

      try {
        for await (const event of streamAgui(
          AGENT_URL,
          input,
          controller.signal,
        )) {
          switch (event.type) {
            case "TEXT_MESSAGE_START": {
              lastMessageId = event.messageId;
              currentText = "";
              const msgId = ensureAssistantMsg();
              const textBlockId = `text_${event.messageId}`;
              currentTextBlockId = textBlockId;
              const textBlock: TextBlock = {
                id: textBlockId,
                type: "text",
                text: "",
              };
              addBlockToMessage(msgId, textBlock);
              break;
            }

            case "TEXT_MESSAGE_CONTENT": {
              currentText += event.delta;
              if (assistantMsgId && currentTextBlockId) {
                updateBlockInMessage<TextBlock>(
                  assistantMsgId,
                  currentTextBlockId,
                  (b) => ({
                    ...b,
                    text: b.text + event.delta,
                  }),
                );
              }
              break;
            }

            case "TEXT_MESSAGE_END":
              if (assistantMsgId && currentTextBlockId && !currentText.trim()) {
                updateMessage(assistantMsgId, (m) => ({
                  ...m,
                  blocks: m.blocks.filter((b) => b.id !== currentTextBlockId),
                }));
              } else {
                completedTextBlockId = currentTextBlockId;
              }
              currentTextBlockId = null;
              break;

            case "TOOL_CALL_START": {
              if (toolCallMap.has(event.toolCallId)) break;
              const msgId = ensureAssistantMsg();
              const blockId = `tool_${event.toolCallId}`;
              const parentId = event.parentToolCallId;
              toolCallMap.set(event.toolCallId, {
                blockId,
                argsJson: "",
                name: event.toolCallName,
                parentToolCallId: parentId,
              });
              if (isAgentTool(event.toolCallName)) {
                activeAgentToolCallId = event.toolCallId;
                activeAgentBlockId = blockId;
              }
              const toolBlock: ToolCallBlock = {
                id: blockId,
                type: "tool_call",
                toolCallId: event.toolCallId,
                toolName: event.toolCallName,
                toolArgs: "",
                toolArgsParsed: null,
                toolStatus: "streaming",
                toolResult: null,
                toolResultRaw: "",
              };
              if (parentId && activeAgentBlockId) {
                updateBlockInMessage<ToolCallBlock>(
                  msgId,
                  activeAgentBlockId,
                  (b) => ({
                    ...b,
                    subAgentToolCalls: [...(b.subAgentToolCalls || []), toolBlock],
                  }),
                );
              } else {
                addBlockToMessage(msgId, toolBlock);
              }
              break;
            }

            case "TOOL_CALL_ARGS": {
              const tc = toolCallMap.get(event.toolCallId);
              if (tc && assistantMsgId) {
                tc.argsJson += event.delta;
                updateToolBlock(tc, (b) => ({ ...b, toolArgs: b.toolArgs + event.delta }));
              }
              break;
            }

            case "TOOL_CALL_END": {
              const tc = toolCallMap.get(event.toolCallId);
              if (tc && assistantMsgId) {
                let parsed: Record<string, unknown> | null = null;
                try {
                  parsed = JSON.parse(tc.argsJson);
                } catch {
                  // partial JSON
                }
                updateToolBlock(tc, (b) => ({ ...b, toolStatus: "executing" as const, toolArgsParsed: parsed }));
                if (!tc.parentToolCallId) {
                  toolCallsForHistory.push({
                    id: event.toolCallId,
                    type: "function",
                    function: { name: tc.name, arguments: tc.argsJson },
                  });
                }
              }
              break;
            }

            case "TOOL_CALL_RESULT": {
              const tc = toolCallMap.get(event.toolCallId);
              if (tc && assistantMsgId) {
                const parsed = parseToolResult(event.content);
                updateToolBlock(tc, (b) => ({
                  ...b,
                  toolStatus: (parsed.returnCode !== 0 ? "error" : "done") as ToolCallBlock["toolStatus"],
                  toolResult: parsed,
                  toolResultRaw: event.content,
                }));
                if (event.toolCallId === activeAgentToolCallId) {
                  activeAgentToolCallId = null;
                  activeAgentBlockId = null;
                }
                if (!tc.parentToolCallId) {
                  toolResultsForHistory.push({
                    id: `res_${event.toolCallId}`,
                    role: "tool",
                    content: event.content,
                    toolCallId: event.toolCallId,
                  });
                }
              } else if (activeAgentBlockId && assistantMsgId) {
                const orphanParsed = parseToolResult(event.content);
                const orphanBlock: ToolCallBlock = {
                  id: `tool_${event.toolCallId}`,
                  type: "tool_call",
                  toolCallId: event.toolCallId,
                  toolName: "unknown",
                  toolArgs: "",
                  toolArgsParsed: null,
                  toolStatus: orphanParsed.returnCode !== 0 ? "error" : "done",
                  toolResult: orphanParsed,
                  toolResultRaw: event.content,
                };
                updateBlockInMessage<ToolCallBlock>(
                  assistantMsgId,
                  activeAgentBlockId,
                  (b) => ({
                    ...b,
                    subAgentToolCalls: [...(b.subAgentToolCalls || []), orphanBlock],
                  }),
                );
              }
              break;
            }

            case "CUSTOM":
              if (event.name === "tool_progress") {
                const progress = event.value as {
                  toolCallId: string;
                  parentToolCallId: string | null;
                  elapsedSeconds: number;
                };
                const targetBlockId = progress.parentToolCallId
                  ? toolCallMap.get(progress.parentToolCallId)?.blockId
                  : toolCallMap.get(progress.toolCallId)?.blockId;
                if (targetBlockId && assistantMsgId) {
                  updateBlockInMessage<ToolCallBlock>(
                    assistantMsgId,
                    targetBlockId,
                    (b) => ({ ...b, elapsedSeconds: progress.elapsedSeconds }),
                  );
                }
              } else if (event.name === "on_interrupt") {
                const interrupt = event.value as unknown as InterruptPayload;
                setPendingInterrupt(interrupt);
                const msgId = ensureAssistantMsg();
                addBlockToMessage(msgId, {
                  id: `interrupt_${crypto.randomUUID()}`,
                  type: "interrupt",
                  interrupt,
                });
              }
              break;

            case "STATE_SNAPSHOT": {
              const snapshot = event.snapshot as unknown as AnalysisStateSnapshot;
              setAnalysisState(snapshot);
              break;
            }

            case "RUN_ERROR": {
              const message = normalizeDisplayedAgentError(event.message);
              setError(message);
              const msgId = ensureAssistantMsg();
              addBlockToMessage(msgId, {
                id: `error_${crypto.randomUUID()}`,
                type: "error",
                message,
              });
              break;
            }

            case "RUN_FINISHED":
              break;
          }
        }
      } catch (err: unknown) {
        if ((err as Error).name !== "AbortError") {
          const msg = normalizeDisplayedAgentError(
            err instanceof Error ? err.message : "Unknown stream error",
          );
          setError(msg);
          const msgId = ensureAssistantMsg();
          addBlockToMessage(msgId, {
            id: `error_${crypto.randomUUID()}`,
            type: "error",
            message: msg,
          });
        }
      } finally {
        setIsStreaming(false);
        abortRef.current = null;
      }

      if (currentText || toolCallsForHistory.length > 0) {
        const finalText = currentText
          ? await normalizeAssistantMarkdownForRender(currentText)
          : currentText;

        if (
          finalText !== currentText &&
          assistantMsgId &&
          completedTextBlockId
        ) {
          updateBlockInMessage<TextBlock>(
            assistantMsgId,
            completedTextBlockId,
            (block) => ({
              ...block,
              text: finalText,
            }),
          );
        }

        const assistantMsg: AguiMessage = {
          id: lastMessageId || crypto.randomUUID(),
          role: "assistant",
          content: finalText,
          ...(toolCallsForHistory.length > 0
            ? { toolCalls: toolCallsForHistory }
            : {}),
        };
        historyRef.current.push(assistantMsg, ...toolResultsForHistory);
        onMessagesChangeRef.current?.(
          messagesRef.current,
          historyRef.current,
          analysisStateRef.current,
          pendingInterruptRef.current,
        );
      }
    },
    [addMessage, addBlockToMessage, updateBlockInMessage, updateMessage],
  );

  const compactHistoryIfNeeded = useCallback(async () => {
    const history = historyRef.current;
    if (
      !history ||
      history.length < publicConfig.agentMinHistoryMessagesForCompaction ||
      estimateHistoryTokens(history) < publicConfig.agentCompactionThresholdTokens
    ) {
      return;
    }

    try {
      const response = await fetch(COMPACT_URL, {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
        },
        body: JSON.stringify({
          messages: history,
        }),
      });

      if (!response.ok) {
        throw new Error(`Compaction returned ${response.status}`);
      }

      const payload = (await response.json()) as {
        history?: AguiMessage[];
      };

      if (Array.isArray(payload.history) && payload.history.length > 0) {
        historyRef.current = payload.history;
        onMessagesChangeRef.current?.(
          messagesRef.current,
          historyRef.current,
          analysisStateRef.current,
          pendingInterruptRef.current,
        );
      }
    } catch (err) {
      console.error("history compaction failed:", err);
    }
  }, []);

  const sendMessage = useCallback(
    async (text: string) => {
      if (!text.trim() || isStreaming) return;

      await compactHistoryIfNeeded();

      addMessage({
        id: crypto.randomUUID(),
        role: "user",
        blocks: [
          { id: crypto.randomUUID(), type: "text", text },
        ],
        timestamp: Date.now(),
      });

      const aguiUserMsg: AguiMessage = {
        id: crypto.randomUUID(),
        role: "user",
        content: text,
      };
      historyRef.current.push(aguiUserMsg);

      const input: RunAgentInput = {
        threadId: threadIdRef.current,
        runId: crypto.randomUUID(),
        messages: historyRef.current,
        tools: [],
        context: [],
        state: {},
        forwardedProps: {},
        repositoryId: options.repositoryId || undefined,
      };

      await processStream(input);
    },
    [isStreaming, processStream, addMessage, compactHistoryIfNeeded],
  );

  const resolveInterrupt = useCallback(
    async (approved: boolean, reason?: string) => {
      if (!pendingInterrupt) return;
      setPendingInterrupt(null);

      const input: RunAgentInput = {
        threadId: threadIdRef.current,
        runId: crypto.randomUUID(),
        messages: historyRef.current,
        tools: [],
        context: [],
        state: {},
        forwardedProps: {
          command: {
            resume: JSON.stringify({
              approved,
              ...(reason ? { reason } : {}),
            }),
          },
        },
      };

      await processStream(input);
    },
    [pendingInterrupt, processStream],
  );

  const resetSession = useCallback(
    (opts: {
      messages: ChatMessage[];
      history: AguiMessage[];
      threadId: string;
      analysisState?: AnalysisStateSnapshot | null;
      pendingInterrupt?: InterruptPayload | null;
    }) => {
      abortRef.current?.abort();
      abortRef.current = null;

      setMessages(opts.messages || []);
      historyRef.current = opts.history || [];
      threadIdRef.current = opts.threadId;
      setError(null);
      setAnalysisState(opts.analysisState ?? null);
      setPendingInterrupt(opts.pendingInterrupt ?? null);
      setIsStreaming(false);
    },
    [],
  );

  const onMessagesChangeRef = useRef(options.onMessagesChange);
  const analysisStateRef = useRef<AnalysisStateSnapshot | null>(null);
  const pendingInterruptRef = useRef<InterruptPayload | null>(null);
  onMessagesChangeRef.current = options.onMessagesChange;

  useEffect(() => {
    messagesRef.current = messages;
  }, [messages]);

  useEffect(() => {
    analysisStateRef.current = analysisState;
  }, [analysisState]);

  useEffect(() => {
    pendingInterruptRef.current = pendingInterrupt;
  }, [pendingInterrupt]);

  useEffect(() => {
    onMessagesChangeRef.current?.(
      messages,
      historyRef.current,
      analysisState,
      pendingInterrupt,
    );
  }, [messages, analysisState, pendingInterrupt]);

  return {
    messages,
    isStreaming,
    error,
    pendingInterrupt,
    analysisState,
    threadId: threadIdRef.current,
    sendMessage,
    resolveInterrupt,
    resetSession,
  };
}

function estimateHistoryTokens(messages: AguiMessage[]): number {
  const totalChars = messages.reduce((sum, message) => {
    const toolChars = (message.toolCalls || []).reduce(
      (toolSum, toolCall) =>
        toolSum +
        toolCall.function.name.length +
        (toolCall.function.arguments?.length || 0),
      0,
    );

    return sum + (message.content?.length || 0) + toolChars;
  }, 0);

  return Math.ceil(totalChars / 4);
}

function parseToolResult(raw: string): {
  stdout: string;
  stderr: string;
  returnCode: number;
} {
  try {
    const parsed = JSON.parse(raw);
    return {
      stdout: parsed.stdout ?? parsed.content ?? "",
      stderr: parsed.stderr ?? "",
      returnCode: parsed.return_code ?? parsed.returnCode ?? 0,
    };
  } catch {
    return { stdout: raw, stderr: "", returnCode: 0 };
  }
}