// dashboard/src/features/reader/AIAssistantPanel.tsx
// 文献 AI 问答助手 — 使用共享 AgentMarkdown / useAgentSSE / 子代理路由 / 渲染组件
import { useState, useRef, useEffect } from 'react';
import {
  Send, Loader, X, BookOpen, AlertCircle, Compass,
  Brain, Square, Paperclip,
} from 'lucide-react';
import axios from 'axios';
import {
  AgentMarkdown, useAgentSSE, useAutoScroll,
  ThoughtCard, ToolCallCard, SubAgentContainer,
  PARENT_COLORS,
  routeThought, routeToolCall, routeToolResult, routeTextDelta,
} from './agent';
import type { SSEEventHandlers, TimelineItem } from './agent';

interface RetrievalSource {
  bibcode: string;
  paragraph_index: number;
  content: string;
  distance: number;
}

interface Message {
  sender: 'user' | 'ai';
  text: string;
  sources?: RetrievalSource[];
  imageUrl?: string;
  timeline: TimelineItem[];
}

interface AIAssistantPanelProps {
  bibcode: string;
  onJumpToSource: (bibcode: string, paragraphIndex: number) => void;
}

const SUGGESTED_QUESTIONS = [
  '对比我馆藏的文献中，针对热亚矮星双星在共同包层抛射过程中恒星风流失速率的各种主流观点差异。',
  '有哪些文献提及了脉动白矮星的非径向振动模？',
  '简述目前文献中关于 Gaia DR3 视差零点改正的处理方法。',
  '文献库中关于双星合并前奏（Precursor）观测特征 of 论述有哪些？',
];

export function AIAssistantPanel({
  bibcode,
  onJumpToSource,
}: AIAssistantPanelProps) {
  const [messages, setMessages] = useState<Message[]>([]);
  const [input, setInput] = useState('');
  const [error, setError] = useState<string | null>(null);
  const [sessionId, setSessionId] = useState<string | null>(null);

  // 折叠状态（共享 ThoughtCard / ToolCallCard / SubAgentContainer 的展开控制）
  const [expandedThoughts, setExpandedThoughts] = useState<Record<string, boolean>>({});
  const [expandedArgs, setExpandedArgs] = useState<Record<string, boolean>>({});
  const [expandedResults, setExpandedResults] = useState<Record<string, boolean>>({});
  const [collapsedSubAgents, setCollapsedSubAgents] = useState<Record<string, boolean>>({});

  const [pendingImage, setPendingImage] = useState<{
    data?: string; path?: string; mime_type: string; name: string;
  } | null>(null);
  const fileInputRef = useRef<HTMLInputElement>(null);

  const { streaming, send, stop: sseStop } = useAgentSSE();
  const { chatEndRef, scrollContainerRef, setShouldAutoScroll, handleScroll } =
    useAutoScroll([messages, streaming]);

  // ── 折叠控制 ──
  const toggleThought = (key: string) =>
    setExpandedThoughts(prev => ({ ...prev, [key]: !prev[key] }));
  const toggleArgs = (tcId: string) =>
    setExpandedArgs(prev => ({ ...prev, [tcId]: !prev[tcId] }));
  const toggleResult = (tcId: string) =>
    setExpandedResults(prev => ({ ...prev, [tcId]: !prev[tcId] }));
  const toggleSubAgent = (id: string) =>
    setCollapsedSubAgents(prev => ({ ...prev, [id]: !prev[id] }));

  // ── 图片处理 ──
  const attachImage = (file: File) => {
    if (!file.type.startsWith('image/')) return;
    const reader = new FileReader();
    reader.onload = (e) => {
      const result = e.target?.result as string;
      const commaIdx = result.indexOf(',');
      if (commaIdx === -1) return;
      setPendingImage({
        data: result.substring(commaIdx + 1),
        mime_type: file.type,
        name: file.name,
      });
    };
    reader.readAsDataURL(file);
  };

  const handlePaste = (e: React.ClipboardEvent) => {
    const items = e.clipboardData?.items;
    if (!items) return;
    for (let i = 0; i < items.length; i++) {
      if (items[i].type.startsWith('image/')) {
        e.preventDefault();
        const file = items[i].getAsFile();
        if (file) attachImage(file);
        return;
      }
    }
  };

  // ── bibcode 切换重置 ──
  useEffect(() => {
    setMessages([]);
    setSessionId(null);
    setError(null);
    setExpandedThoughts({});
    setExpandedArgs({});
    setExpandedResults({});
    setCollapsedSubAgents({});
    setPendingImage(null);
  }, [bibcode]);

  const handleStop = async () => {
    sseStop();
    if (sessionId) {
      try { await axios.post(`/api/chat/sessions/${sessionId}/stop`); } catch { /* noop */ }
    }
  };

  const handleSend = async (questionText: string) => {
    if ((!questionText.trim() && !pendingImage) || streaming) return;

    const figure = pendingImage;
    setError(null);
    setPendingImage(null);

    const userMsg: Message = {
      sender: 'user',
      text: questionText,
      timeline: [],
      imageUrl: figure
        ? (figure.path
            ? `/api/files/${figure.path}`
            : `data:${figure.mime_type};base64,${figure.data}`)
        : undefined,
    };
    setMessages(prev => [...prev, userMsg]);
    setInput('');
    setShouldAutoScroll(true);

    // AI 消息占位符（用 TimelineItem[] 替代旧 AgentStep[]）
    const aiMsg: Message = { sender: 'ai', text: '', timeline: [], sources: [] };
    setMessages(prev => [...prev, aiMsg]);

    // ── SSE 事件处理器（使用共享子代理路由）──
    const handlers: SSEEventHandlers = {
      onSession: (sid) => {
        if (!sessionId) setSessionId(sid);
      },
      onThought: (step, content) => {
        setMessages(prev => updateAiTimeline(prev, (timeline) =>
          routeThought(timeline, step, content),
        ));
      },
      onToolCall: (step, _id, name, args) => {
        setMessages(prev => updateAiTimeline(prev, (timeline) =>
          routeToolCall(timeline, step, _id, name, args),
        ));
      },
      onToolResult: (_step, tcId, name, _output, isError, metadata) => {
        setExpandedResults(prev => ({ ...prev, [tcId]: true }));
        setMessages(prev => {
          if (prev.length === 0) return prev;
          const last = prev[prev.length - 1];
          if (last.sender !== 'ai') return prev;
          const newTimeline = routeToolResult(last.timeline, tcId, name, _output, isError);
          // RAG 来源收集
          let sources = [...(last.sources || [])];
          if (name === 'rag_search' && metadata?.sources) {
            for (const s of metadata.sources) {
              if (!sources.some(
                c => c.bibcode === s.bibcode && c.paragraph_index === s.paragraph_index,
              )) {
                sources.push({
                  bibcode: s.bibcode, paragraph_index: s.paragraph_index,
                  content: s.preview || '', distance: s.distance,
                });
              }
            }
          }
          return [...prev.slice(0, -1), { ...last, timeline: newTimeline, sources }];
        });
      },
      onTextDelta: (content, toolCallId) => {
        if (toolCallId) {
          setExpandedResults(prev => ({ ...prev, [toolCallId]: true }));
        }
        setMessages(prev => {
          const last = prev[prev.length - 1];
          if (last?.sender !== 'ai') return prev;
          const { timeline, answerDelta } = routeTextDelta(last.timeline, content, toolCallId);
          return [
            ...prev.slice(0, -1),
            { ...last, timeline, text: last.text + answerDelta },
          ];
        });
      },
      onError: (msg) => {
        setError(msg);
      },
    };

    const contextPrefix = `[当前正在阅读文献: ${bibcode}] `;
    const requestQuery = questionText.includes(bibcode)
      ? questionText : `${contextPrefix}${questionText}`;

    let imagePayload: { data: string; mime_type: string } | undefined;
    if (figure?.data) {
      imagePayload = { data: figure.data, mime_type: figure.mime_type };
    }

    await send(
      {
        question: requestQuery,
        sessionId,
        mode: 'literature-reader',
        image: imagePayload
          ? { data: imagePayload.data, mime_type: imagePayload.mime_type }
          : undefined,
      },
      handlers,
    );
  };

  // ── 时间线条目渲染（委托到共享组件）──
  const renderTimelineItem = (item: TimelineItem, idx: number, isStreaming: boolean) => {
    switch (item.type) {
      case 'subagent_container':
        return (
          <SubAgentContainer
            key={`subagent-${item.id}-${idx}`}
            status={item.status}
            children={item.children}
            summary={item.summary}
            isStreaming={isStreaming}
            isCollapsed={collapsedSubAgents[item.id] !== false}
            onToggle={() => toggleSubAgent(item.id)}
            expandedThoughts={expandedThoughts}
            expandedArgs={expandedArgs}
            expandedResults={expandedResults}
            onToggleThought={toggleThought}
            onToggleArgs={toggleArgs}
            onToggleResult={toggleResult}
          />
        );
      case 'thought': {
        const thoughtKey = `ai-thought-${item.step}-${idx}`;
        return (
          <ThoughtCard
            key={thoughtKey}
            step={item.step}
            content={item.content}
            isStreaming={isStreaming}
            colorScheme={PARENT_COLORS}
            isExpanded={expandedThoughts[thoughtKey] === true}
            onToggle={() => toggleThought(thoughtKey)}
          />
        );
      }
      case 'tool_call': {
        const tcId = item.id;
        return (
          <ToolCallCard
            key={`tc-${tcId}-${idx}`}
            step={item.step}
            name={item.name}
            arguments={item.arguments}
            result={item.result}
            isStreaming={isStreaming}
            colorScheme={PARENT_COLORS}
            isArgsExpanded={expandedArgs[tcId] === true}
            isResultExpanded={expandedResults[tcId] === true}
            onToggleArgs={() => toggleArgs(tcId)}
            onToggleResult={() => toggleResult(tcId)}
          />
        );
      }
      case 'answer':
        return null; // finalAnswer 单独渲染在气泡内
      default:
        return null;
    }
  };

  return (
    <div className="flex flex-col overflow-hidden h-full">
      {/* 消息区 */}
      <div
        ref={scrollContainerRef}
        onScroll={handleScroll}
        className="flex-1 overflow-y-auto p-4 pb-28 space-y-4 min-h-0"
      >
        {messages.length === 0 ? (
          <div className="py-6 space-y-6">
            <div className="text-center space-y-2 max-w-sm mx-auto">
              <Compass className="w-10 h-10 mx-auto text-sky-500 opacity-60" />
              <h3 className="text-xs font-bold text-slate-800">文献阅读助手</h3>
              <p className="text-[11px] text-slate-500 leading-relaxed font-semibold">
                专注文献精读与理解的 AI 助手。支持语义检索、逐段解读、公式分析，严格只读模式保障文献安全。
              </p>
            </div>
            <div className="space-y-2">
              <span className="text-[10px] font-bold text-slate-400 tracking-wider uppercase">
                推荐学术提问:
              </span>
              <div className="space-y-2">
                {SUGGESTED_QUESTIONS.map((q, idx) => (
                  <button
                    key={idx}
                    onClick={() => handleSend(q)}
                    className="w-full text-left bg-white hover:bg-slate-55 border border-slate-200 rounded-md p-3 text-xs text-slate-700 leading-relaxed font-medium transition-all cursor-pointer hover:border-slate-350"
                  >
                    {q}
                  </button>
                ))}
              </div>
            </div>
          </div>
        ) : (
          messages.map((msg, index) => (
            <div
              key={index}
              className={`flex flex-col ${msg.sender === 'user' ? 'items-end' : 'items-start'} space-y-1.5`}
            >
              <span className="text-[10px] font-bold text-slate-400 px-1">
                {msg.sender === 'user' ? '我' : '天文科研助手'}
              </span>
              <div
                className={`max-w-[95%] rounded-lg px-4 py-3 text-xs leading-relaxed font-medium shadow-xs border ${
                  msg.sender === 'user'
                    ? 'bg-blueprint text-white border-blueprint select-text'
                    : 'bg-white text-slate-800 border-slate-250 select-text'
                }`}
              >
                {msg.sender === 'user' ? (
                  <div className="space-y-2">
                    {msg.imageUrl && (
                      <img
                        src={msg.imageUrl}
                        alt="Attached Figure"
                        className="max-w-40 max-h-40 object-contain rounded-md border border-sky-500/30"
                      />
                    )}
                    <p className="whitespace-pre-wrap">{msg.text}</p>
                  </div>
                ) : (
                  <div className="space-y-2.5">
                    {/* 时间线条目（思考/工具调用/子代理容器）*/}
                    {msg.timeline.length > 0 && (
                      <div className="border border-slate-205 bg-slate-100/30 rounded-md p-2 select-none">
                        <div className="text-[10px] font-bold text-slate-600 flex items-center gap-1.5 mb-2 px-1">
                          <Brain className="w-3.5 h-3.5 text-slate-450" />
                          <span>推理与工具链 ({msg.timeline.length} 项)</span>
                        </div>
                        <div className="pl-4 border-l border-slate-200 space-y-2">
                          {msg.timeline.map((item, tIdx) =>
                            renderTimelineItem(item, tIdx, !!streaming),
                          )}
                        </div>
                      </div>
                    )}

                    {/* 最终回答 */}
                    {msg.text ? (
                      <div className="prose prose-sm max-w-none text-slate-800 leading-relaxed prose-headings:text-slate-900 prose-headings:font-bold prose-strong:text-slate-900 prose-code:text-sky-700 prose-img:rounded-lg">
                        <AgentMarkdown>{msg.text}</AgentMarkdown>
                      </div>
                    ) : (
                      !streaming && (
                        <span className="text-slate-400 italic">正在构建最终结论...</span>
                      )
                    )}
                  </div>
                )}
              </div>

              {/* 关联来源 */}
              {msg.sender === 'ai' && msg.sources && msg.sources.length > 0 && (
                <div className="w-full mt-2 pl-2 space-y-1.5 border-l-2 border-slate-200">
                  <span className="text-[9px] font-bold text-slate-400 select-none">
                    参考来源文献（点击跳转）:
                  </span>
                  <div className="grid grid-cols-1 gap-1.5 w-[95%]">
                    {msg.sources.map((src, sIdx) => (
                      <button
                        key={sIdx}
                        onClick={() => onJumpToSource(src.bibcode, src.paragraph_index)}
                        className="flex items-center justify-between text-left bg-white hover:bg-slate-100 border border-slate-200 rounded-md px-2.5 py-1.5 text-[10px] text-slate-655 font-semibold transition-all shadow-sm cursor-pointer hover:border-sky-300"
                        title={src.content}
                      >
                        <div className="flex items-center gap-1.5 truncate mr-2">
                          <BookOpen className="w-3 h-3 text-sky-600 shrink-0" />
                          <span className="font-bold text-slate-700 truncate">{src.bibcode}</span>
                          <span className="text-slate-400 text-[9px] font-bold">
                            §{src.paragraph_index + 1}
                          </span>
                        </div>
                        <span className="text-slate-400 text-[9px] shrink-0 font-medium">
                          距: {src.distance.toFixed(3)}
                        </span>
                      </button>
                    ))}
                  </div>
                </div>
              )}
            </div>
          ))
        )}

        {streaming && messages.length > 0 && !messages[messages.length - 1].text && (
          <div className="flex items-center space-x-2 text-slate-500 pl-2">
            <Loader className="w-4 h-4 animate-spin text-sky-600" />
            <span className="text-[10px] font-bold">文献检索及推理结论构建中...</span>
          </div>
        )}

        {error && (
          <div className="flex items-start gap-2 bg-red-50 border border-red-200 text-red-700 rounded-lg p-3 text-xs w-[95%] font-semibold">
            <AlertCircle className="w-4 h-4 text-red-500 shrink-0 mt-0.5" />
            <div>
              <div className="font-bold">查询发生错误</div>
              <div className="text-[10px] opacity-80 mt-0.5">{error}</div>
            </div>
          </div>
        )}

        <div ref={chatEndRef} />
      </div>

      {/* 底栏输入 */}
      <form
        onSubmit={(e) => {
          e.preventDefault();
          handleSend(input);
        }}
        className="absolute bottom-0 left-0 right-0 p-3 bg-transparent pointer-events-none shrink-0"
      >
        <div className="bg-slate-50 border border-slate-300 rounded-lg p-2.5 focus-within:bg-white focus-within:border-blueprint focus-within:ring-1 focus-within:ring-blueprint/10 transition-all relative pointer-events-auto shadow-sm flex flex-col gap-2">
          <textarea
            value={input}
            onChange={(e) => setInput(e.target.value)}
            onKeyDown={(e) => {
              if (e.key === 'Enter' && !e.shiftKey) {
                e.preventDefault();
                handleSend(input);
              }
            }}
            onPaste={handlePaste}
            disabled={streaming}
            placeholder={pendingImage ? '针对选中图表提问...' : '向 AI 馆藏助手提问...'}
            rows={2}
            className="w-full bg-transparent resize-none border-none outline-none focus:outline-none focus:ring-0 text-xs text-slate-900 placeholder-slate-400 leading-relaxed font-semibold min-h-[36px] max-h-32 pb-1"
          />
          <div className="flex justify-between items-center mt-1.5 pt-1.5 border-t border-slate-200/40">
            <div className="flex gap-2 items-center min-w-0">
              <input
                ref={fileInputRef}
                type="file"
                accept="image/*"
                className="hidden"
                onChange={(e) => {
                  const file = e.target.files?.[0];
                  if (file) attachImage(file);
                  e.target.value = '';
                }}
              />
              <button
                type="button"
                onClick={() => fileInputRef.current?.click()}
                disabled={streaming}
                title="上传或粘贴图片（也可直接 Ctrl+V 粘贴）"
                className={`p-1 rounded-md border text-[10px] font-bold transition-all cursor-pointer flex items-center justify-center shrink-0 ${
                  pendingImage
                    ? 'bg-sky-50 border-sky-300 text-sky-700'
                    : 'bg-white border-slate-200 text-slate-400 hover:text-sky-600 hover:border-sky-200'
                } disabled:opacity-60`}
              >
                <Paperclip className="w-3.5 h-3.5" />
              </button>
              {pendingImage && (
                <div className="relative inline-block ml-1 shrink-0 group select-none">
                  <img
                    src={
                      pendingImage.path
                        ? `/api/files/${pendingImage.path}`
                        : `data:${pendingImage.mime_type};base64,${pendingImage.data}`
                    }
                    alt="Preview"
                    className="w-5 h-5 object-cover rounded border border-slate-300 shadow-2xs"
                  />
                  <button
                    type="button"
                    onClick={() => setPendingImage(null)}
                    className="absolute -top-1 -right-1 bg-slate-500/80 hover:bg-slate-700 text-white rounded-full w-3 h-3 shadow-xs transition-colors cursor-pointer flex items-center justify-center"
                    title="移除图片"
                  >
                    <X className="w-2 h-2" />
                  </button>
                </div>
              )}
            </div>
            <div>
              {streaming ? (
                <button
                  type="button"
                  onClick={handleStop}
                  className="p-1 rounded-lg bg-slate-800 hover:bg-slate-950 text-white transition-colors cursor-pointer flex items-center justify-center shadow-xs"
                  title="手动停止执行"
                >
                  <Square className="w-3.5 h-3.5" />
                </button>
              ) : (
                <button
                  type="submit"
                  disabled={!input.trim() && !pendingImage}
                  className="p-1 px-2.5 rounded-lg bg-sky-600 hover:bg-sky-700 text-white disabled:bg-slate-100 disabled:text-slate-350 transition-colors cursor-pointer flex items-center justify-center gap-1 shadow-xs text-[10px] font-bold"
                >
                  <Send className="w-3.5 h-3.5" />
                  <span>发送</span>
                </button>
              )}
            </div>
          </div>
        </div>
      </form>
    </div>
  );
}

// ── helpers ──

// 更新最后一条 AI 消息的 timeline
function updateAiTimeline(
  prev: Message[],
  fn: (timeline: TimelineItem[]) => TimelineItem[],
): Message[] {
  if (prev.length === 0) return prev;
  const last = prev[prev.length - 1];
  if (last.sender !== 'ai') return prev;
  return [...prev.slice(0, -1), { ...last, timeline: fn(last.timeline) }];
}

