import { useState, useEffect, useCallback, useRef } from 'react';
import { chatSocket } from '@/lib/websocket';
import type { Message } from '@/types';

interface PipelineState {
  phase: string;
  progress: number;
  currentItem?: number;
  totalItems?: number;
  status?: string;
}

export function useChat(conversationId: string | null) {
  const [messages, setMessages] = useState<Message[]>([]);
  const [connected, setConnected] = useState(false);
  const [thinking, setThinking] = useState(false);
  const [toolStatus, setToolStatus] = useState<string | null>(null);
  const [pipeline, setPipeline] = useState<PipelineState | null>(null);

  // Track which conversation is currently loaded via WS
  const activeConvRef = useRef<string | null>(null);
  // Pending conversation creation resolve
  const createResolveRef = useRef<((id: string) => void) | null>(null);

  useEffect(() => {
    const onConnected = () => setConnected(true);
    const onHistory = (data: Record<string, unknown>) => {
      setMessages((data.messages as Message[]) || []);
    };
    const onMessage = (data: Record<string, unknown>) => {
      setMessages((prev) => [...prev, data as unknown as Message]);
    };
    const onStatus = (data: Record<string, unknown>) => {
      if (data.status === 'thinking') setThinking(true);
      if (data.status === 'done') { setThinking(false); setToolStatus(null); }
    };
    const onMessageStart = (data: Record<string, unknown>) => {
      setThinking(false);
      setMessages((prev) => [...prev, {
        id: data.id as string,
        role: 'assistant' as const,
        content: '',
        reasoningContent: '',
        created_at: new Date().toISOString(),
        conversation_id: '',
      }]);
    };
    const onReasoningDelta = (data: Record<string, unknown>) => {
      setMessages((prev) => prev.map((m) =>
        m.id === data.id ? { ...m, reasoningContent: (m.reasoningContent || '') + (data.text as string) } : m
      ));
    };
    const onTextDelta = (data: Record<string, unknown>) => {
      setMessages((prev) => prev.map((m) =>
        m.id === data.id ? { ...m, content: m.content + (data.text as string) } : m
      ));
    };
    const onMessageEnd = () => {
      setThinking(false);
    };
    const onToolStart = (data: Record<string, unknown>) => {
      setToolStatus(`执行: ${data.tool}...`);
    };
    const onToolResult = (data: Record<string, unknown>) => {
      setToolStatus(null);
      setThinking(true);
      setMessages((prev) => [...prev, {
        id: `tool-${Date.now()}`,
        role: 'tool' as const,
        content: JSON.stringify({ tool: data.tool, result: data.result }),
        assets: (data.assets as Message['assets']) || undefined,
        created_at: new Date().toISOString(),
        conversation_id: '',
      }]);
    };
    const onToolError = (data: Record<string, unknown>) => {
      setToolStatus(`失败: ${data.tool}`);
      setTimeout(() => setToolStatus(null), 4000);
    };
    const onPipelineProgress = (data: Record<string, unknown>) => {
      setPipeline({
        phase: data.phase as string,
        progress: data.progress as number,
        currentItem: data.currentItem as number | undefined,
        totalItems: data.totalItems as number | undefined,
        status: data.status as string | undefined,
      });
    };
    const onConversationCreated = (data: Record<string, unknown>) => {
      const id = data.id as string;
      activeConvRef.current = id;
      if (createResolveRef.current) {
        createResolveRef.current(id);
        createResolveRef.current = null;
      }
    };

    chatSocket.on('connected', onConnected);
    chatSocket.on('history', onHistory);
    chatSocket.on('message', onMessage);
    chatSocket.on('status', onStatus);
    chatSocket.on('message_start', onMessageStart);
    chatSocket.on('reasoning_delta', onReasoningDelta);
    chatSocket.on('text_delta', onTextDelta);
    chatSocket.on('message_end', onMessageEnd);
    chatSocket.on('tool_start', onToolStart);
    chatSocket.on('tool_result', onToolResult);
    chatSocket.on('tool_error', onToolError);
    chatSocket.on('pipeline_progress', onPipelineProgress);
    chatSocket.on('conversation_created', onConversationCreated);
    chatSocket.connect();

    return () => {
      chatSocket.off('connected', onConnected);
      chatSocket.off('history', onHistory);
      chatSocket.off('message', onMessage);
      chatSocket.off('status', onStatus);
      chatSocket.off('message_start', onMessageStart);
      chatSocket.off('reasoning_delta', onReasoningDelta);
      chatSocket.off('text_delta', onTextDelta);
      chatSocket.off('message_end', onMessageEnd);
      chatSocket.off('tool_start', onToolStart);
      chatSocket.off('tool_result', onToolResult);
      chatSocket.off('tool_error', onToolError);
      chatSocket.off('pipeline_progress', onPipelineProgress);
      chatSocket.off('conversation_created', onConversationCreated);
      chatSocket.disconnect();
    };
  }, []);

  // Init/re-init when conversationId changes
  useEffect(() => {
    if (!connected) return;

    if (conversationId && conversationId !== activeConvRef.current) {
      activeConvRef.current = conversationId;
      setMessages([]);
      setThinking(false);
      setToolStatus(null);
      setPipeline(null);
      chatSocket.send('init', { conversationId });
    }

    if (!conversationId) {
      activeConvRef.current = null;
      setMessages([]);
      setThinking(false);
      setToolStatus(null);
      setPipeline(null);
    }
  }, [conversationId, connected]);

  const send = useCallback((content: string) => {
    chatSocket.send('chat', { content });
  }, []);

  const stop = useCallback(() => {
    chatSocket.stop();
  }, []);

  const createConversation = useCallback((title: string, accountId?: string): Promise<string> => {
    return new Promise((resolve) => {
      createResolveRef.current = resolve;
      chatSocket.send('create_conversation', { title, accountId });
      // Timeout fallback
      setTimeout(() => {
        if (createResolveRef.current) {
          createResolveRef.current = null;
          resolve('');
        }
      }, 5000);
    });
  }, []);

  const removeMessage = useCallback(async (msgId: string) => {
    setMessages((prev) => prev.filter((m) => m.id !== msgId));
    try { await fetch(`/api/pipeline/messages/${msgId}`, { method: 'DELETE' }); } catch {}
  }, []);

  return { messages, connected, thinking, toolStatus, pipeline, send, stop, createConversation, removeMessage };
}