import { useState, useEffect, useRef, useCallback } from 'react';
import { useChatWebSocket } from 'hooks/useChatWebSocket';
import { useClient } from 'hooks/useClient';
import { useSettings } from 'state/useSettings';
import type { ChatMessage, ServerMessage, TaskInfo, Message } from '../apps/Chat/types';

const SAVE_DEBOUNCE_MS = 1000;

type ResourceChatStorage = {
  load: () => Promise<{ sessionId: string | null; messages: ChatMessage[] }>;
  save: (sessionId: string, messages: ChatMessage[]) => Promise<void>;
};

type UsePiChatOptions = {
  replaceUrl?: boolean;
  storage?: ResourceChatStorage;
  resourceChatDir?: string;
  taskInfo?: TaskInfo;
  projectScoped?: boolean;
  context?: string;
  contextId?: string;
  resumeSummary?: string;
  resumeSessionId?: string;
  initialMessages?: ChatMessage[];
  onTurnComplete?: (hadToolCalls: boolean) => void;
};

export function useChat(initialSessionId?: string, initialModel?: string | null, options?: UsePiChatOptions) {
  const {
    replaceUrl = true,
    storage,
    resourceChatDir,
    taskInfo,
    projectScoped,
    context,
    contextId,
    resumeSummary: initialResumeSummary,
    resumeSessionId,
    initialMessages: preloadedMessages,
    onTurnComplete,
  } = options ?? {};
  const [messages, setMessages] = useState<ChatMessage[]>(preloadedMessages ?? []);
  const [streamingText, setStreamingText] = useState('');
  const [isGenerating, setIsGenerating] = useState(false);
  const [sessionId, setSessionId] = useState<string | null>(initialSessionId ?? null);
  const [model, setModel] = useState<string | null>(null);
  const [selectedModel, setSelectedModel] = useState<string | null>(initialModel ?? null);
  const [cwd, setCwd] = useState<string | null>(null);
  const [thinkingLevel, setThinkingLevel] = useState<string | null>(null);

  // Resilient-transport cursor: the max durable-event `seq` seen. On (re)connect we send it as a
  // resume-cursor so the server replays anything we missed (e.g. a background task:notification that
  // landed while briefly disconnected). sendRef breaks the hook ↔ onOpen circular dependency.
  const cursorRef = useRef(0);
  const sendRef = useRef<(d: Record<string, unknown>) => void>(() => {});

  // Track if session has started (first message sent)
  const [hasStarted, setHasStarted] = useState(!!preloadedMessages?.length);

  // Get user settings for default model
  const { settings } = useSettings();

  // Set default model from settings when starting a new chat (no initialSessionId, no initialModel)
  useEffect(() => {
    if (!initialSessionId && !initialModel) {
      const defaultModel = (projectScoped && settings?.chat?.defaultProjectModel) || settings?.chat?.defaultModel;
      if (defaultModel) setSelectedModel(defaultModel);
    }
  }, [initialSessionId, initialModel, settings, projectScoped]);

  // Sync selectedModel when initialModel changes (e.g. resuming a session)
  useEffect(() => {
    if (initialModel) setSelectedModel(initialModel);
  }, [initialModel]);

  const streamingRef = useRef('');
  const rafRef = useRef<number | null>(null);
  const sessionIdRef = useRef<string | null>(initialSessionId ?? null);
  const saveTimerRef = useRef<number | null>(null);
  const toolCallsInTurnRef = useRef(false);
  const onTurnCompleteRef = useRef(onTurnComplete);
  onTurnCompleteRef.current = onTurnComplete;
  const resumeSummaryRef = useRef<string | undefined>(initialResumeSummary);

  const client = useClient();

  const token = localStorage.getItem('BEARER_TOKEN');
  const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
  const wsUrl = `${protocol}//${window.location.host}/api/chat/ws?token=${token}`;

  function flushStreaming() {
    if (rafRef.current !== null) cancelAnimationFrame(rafRef.current);
    rafRef.current = requestAnimationFrame(() => {
      setStreamingText(streamingRef.current);
      rafRef.current = null;
    });
  }

  function commitStreaming() {
    if (!streamingRef.current) return;
    const text = streamingRef.current;
    streamingRef.current = '';
    setStreamingText('');
    setMessages((prev) => [...prev, { role: 'assistant', id: crypto.randomUUID(), text }]);
  }

  function handleMessage(data: unknown) {
    const msg = data as ServerMessage;

    // Advance the resume cursor for any durable (seq-carrying) event.
    const seq = (data as { seq?: number }).seq;
    if (typeof seq === 'number' && seq > cursorRef.current) cursorRef.current = seq;

    switch (msg.type) {
      case 'session:init':
        sessionIdRef.current = msg.sessionId;
        setSessionId(msg.sessionId);
        setModel(msg.model);
        setCwd(msg.cwd);
        if (replaceUrl) window.history.replaceState(null, '', `/chat/${msg.sessionId}`);
        break;

      case 'assistant:delta':
        streamingRef.current += msg.text;
        flushStreaming();
        break;

      case 'assistant:text':
        if (streamingRef.current) {
          commitStreaming();
        } else {
          setMessages((prev) => [...prev, { role: 'assistant', id: crypto.randomUUID(), text: msg.text }]);
        }
        break;

      case 'tool:start':
        toolCallsInTurnRef.current = true;
        setMessages((prev) => [
          ...prev,
          {
            role: 'tool',
            toolName: msg.toolName,
            toolInput: msg.toolInput,
            toolCallId: msg.toolCallId,
          },
        ]);
        break;

      case 'tool:result':
        setMessages((prev) =>
          prev.map((m) =>
            m.role === 'tool' && m.toolCallId === msg.toolCallId
              ? { ...m, output: msg.output, isError: msg.isError }
              : m,
          ),
        );
        break;

      case 'result': {
        commitStreaming();
        setMessages((prev) => [
          ...prev,
          {
            role: 'result',
            cost: msg.cost,
          },
        ]);
        setIsGenerating(false);
        const hadTools = toolCallsInTurnRef.current;
        toolCallsInTurnRef.current = false;
        onTurnCompleteRef.current?.(hadTools);
        break;
      }

      case 'sync:messages':
        sessionIdRef.current = msg.sessionId;
        setSessionId(msg.sessionId);
        // Convert Message[] to ChatMessage[]
        const chatMessages = msg.messages.map((m): ChatMessage => {
          if (m.role === 'user') {
            return { role: 'user', text: m.text || '' };
          } else if (m.role === 'assistant') {
            return { role: 'assistant', id: m.id ?? crypto.randomUUID(), text: m.text || '' };
          } else if (m.role === 'tool') {
            return {
              role: 'tool',
              toolName: m.toolName || '',
              toolInput: m.toolInput || {},
              toolCallId: m.toolCallId || '',
              output: m.output,
              isError: m.isError,
            };
          }
          return { role: 'assistant', text: '' }; // Fallback
        });
        setMessages(chatMessages);
        if (chatMessages.length > 0) setHasStarted(true);
        setIsGenerating(msg.isGenerating);
        if (msg.streamingText) {
          streamingRef.current = msg.streamingText;
          flushStreaming();
        }
        break;

      case 'error':
        commitStreaming();
        setMessages((prev) => [...prev, { role: 'error', text: msg.message }]);
        setIsGenerating(false);
        break;

      case 'stopped':
        commitStreaming();
        setIsGenerating(false);
        break;

      case 'disconnected':
        // Session torn down server-side (turn killed, in-memory session dropped). Commit any partial
        // stream and settle to idle; the WS stays open so a new prompt starts a fresh session.
        commitStreaming();
        setIsGenerating(false);
        break;

      case 'task:started':
        setMessages((prev) => [...prev, { role: 'assistant', id: crypto.randomUUID(), text: `⏳ Background task started — ${msg.description}` }]);
        break;

      case 'task:notification': {
        // The fix in action: a background task's completion arriving after the turn ended.
        const icon = msg.status === 'completed' ? '✅' : msg.status === 'failed' ? '❌' : '⏹️';
        setMessages((prev) => [...prev, { role: 'assistant', id: crypto.randomUUID(), text: `${icon} Background task ${msg.status} — ${msg.summary}` }]);
        break;
      }
    }
  }

  // On every (re)connect, if a session is already established, re-bind + replay via resume-cursor.
  // The first connect (no session yet) no-ops; the first turn establishes the session via session:init.
  const onOpen = useCallback(() => {
    const sid = sessionIdRef.current;
    if (sid) sendRef.current({ type: 'resume-cursor', sessionId: sid, cursor: cursorRef.current });
  }, []);

  const { isConnected, send } = useChatWebSocket({ url: wsUrl, onMessage: handleMessage, onOpen });
  sendRef.current = send;

  // Clean up RAF on unmount
  useEffect(() => {
    return () => {
      if (rafRef.current !== null) cancelAnimationFrame(rafRef.current);
    };
  }, []);

  function sendPrompt(
    text: string,
    attachmentIds?: string[],
    images?: { filename: string; dataUrl: string }[],
    cwdParam?: { root?: string; path: string },
    groupSlug?: string | null,
    thinking?: string | null,
    displayText?: string,
  ) {
    // Mark session as started on first message
    if (!hasStarted) {
      setHasStarted(true);
    }

    setMessages((prev) => [
      ...prev,
      { role: 'user', text: displayText ?? text, ...(images?.length ? { images } : {}) },
    ]);
    setIsGenerating(true);
    streamingRef.current = '';
    setStreamingText('');
    toolCallsInTurnRef.current = false;

    // Parse dataUrls into { mediaType, data } for the server
    const imageData = images
      ?.map((img) => {
        const match = img.dataUrl.match(/^data:([^;]+);base64,(.+)$/);
        return match ? { mediaType: match[1], data: match[2] } : null;
      })
      .filter((x): x is { mediaType: string; data: string } => x !== null);

    const pendingResumeSummary = resumeSummaryRef.current;
    resumeSummaryRef.current = undefined;

    send({
      type: 'chat',
      prompt: text,
      ...(displayText ? { displayText } : {}),
      sessionId: sessionIdRef.current,
      ...(selectedModel ? { model: selectedModel } : {}),
      ...(cwdParam?.path ? { cwd: cwdParam.path } : {}),
      ...(cwdParam?.root ? { cwdRoot: cwdParam.root } : {}),
      ...(groupSlug !== undefined ? { groupSlug } : {}),
      ...(attachmentIds?.length ? { attachmentIds } : {}),
      ...(imageData?.length ? { images: imageData } : {}),
      ...(resourceChatDir ? { resourceChatDir } : {}),
      ...(taskInfo ? { taskInfo } : {}),
      ...(thinking ? { thinking } : {}),
      ...(context ? { context } : {}),
      ...(contextId ? { contextId } : {}),
      ...(pendingResumeSummary ? { resumeSummary: pendingResumeSummary } : {}),
      ...(resumeSessionId ? { resumeSessionId } : {}),
    });
  }

  function stopGeneration() {
    send({ type: 'stop' });
  }

  // Tear down the whole session server-side (kill any in-flight turn + drop the in-memory session),
  // freeing it so its transcript can be resumed elsewhere. Stronger than stopGeneration (turn only).
  function disconnectSession() {
    send({ type: 'disconnect' });
  }

  return {
    messages,
    streamingText,
    isConnected,
    isGenerating,
    sessionId,
    model,
    selectedModel,
    hasStarted,
    cwd,
    thinkingLevel,
    setSelectedModel,
    setThinkingLevel,
    sendPrompt,
    stopGeneration,
    disconnectSession,
  };
}

export type UseChatType = ReturnType<typeof useChat>;
