// dashboard/src/features/reader/AIAssistantPanel.tsx
import { useState, useRef, useEffect } from 'react';
import ReactMarkdown from 'react-markdown';
import remarkMath from 'remark-math';
import remarkGfm from 'remark-gfm';
import rehypeRaw from 'rehype-raw';
import rehypeKatex from 'rehype-katex';
import rehypeSanitize, { defaultSchema } from 'rehype-sanitize';
import 'katex/dist/katex.min.css';
import { 
  Send, Loader, Sparkles, X, BookOpen, AlertCircle, Compass, 
  Brain, Settings, ChevronDown, ChevronUp, AlertTriangle, Square 
} from 'lucide-react';
import axios from 'axios';

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

interface AgentStep {
  step: number;
  type: 'thought' | 'tool_call' | 'tool_result' | 'error';
  label: string;
  detail?: string;
  isError?: boolean;
  isFinished?: boolean;
}

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

interface AIAssistantPanelProps {
  bibcode: string;
  onClose: () => void;
  onJumpToSource: (bibcode: string, paragraphIndex: number) => void;
  pendingFigure?: { path: string; url: string } | null;
  onClearPendingFigure?: () => void;
}

const safeSchema = {
  ...defaultSchema,
  attributes: {
    ...defaultSchema.attributes,
    '*': (defaultSchema.attributes?.['*'] || []).concat(['className', 'style', 'mathvariant', 'display']),
  },
  tagNames: (defaultSchema.tagNames || []).concat([
    'math', 'mrow', 'mi', 'mo', 'mn', 'msup', 'msub', 'msubsup', 'mfrac', 'mover', 'munder', 'munderover', 'mspace', 'mtext', 'annotation'
  ]),
};

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

function getToolDisplayName(name: string): string {
  switch (name) {
    case 'search_papers': return '检索 ADS/arXiv';
    case 'get_paper_metadata': return '获取文献详细元数据';
    case 'download_paper': return '下载文献全文资源';
    case 'parse_paper': return '结构化解析文献内容';
    case 'get_paper_content': return '获取文献全文内容';
    case 'rag_search': return '语义库检索 (RAG)';
    case 'query_target': return '查询天体物理参数';
    default: return name;
  }
}

export function AIAssistantPanel({
  bibcode,
  onClose,
  onJumpToSource,
  pendingFigure,
  onClearPendingFigure
}: AIAssistantPanelProps) {
  const [messages, setMessages] = useState<Message[]>([]);
  const [input, setInput] = useState('');
  const [loading, setLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [sessionId, setSessionId] = useState<string | null>(null);
  
  // 步骤展开状态
  const [stepsExpanded, setStepsExpanded] = useState<Record<number, boolean>>({});

  const chatEndRef = useRef<HTMLDivElement>(null);
  const scrollContainerRef = useRef<HTMLDivElement>(null);
  const [shouldAutoScroll, setShouldAutoScroll] = useState(true);

  const scrollToBottom = () => {
    chatEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  };

  useEffect(() => {
    if (shouldAutoScroll) {
      scrollToBottom();
    }
  }, [messages, loading, shouldAutoScroll]);

  const handleScroll = () => {
    if (!scrollContainerRef.current) return;
    const { scrollTop, scrollHeight, clientHeight } = scrollContainerRef.current;
    // 距离底部 50 像素以内视为贴紧底部
    const isAtBottom = scrollHeight - scrollTop - clientHeight < 50;
    setShouldAutoScroll(isAtBottom);
  };

  // 重置会话（当 bibcode 切换时）
  useEffect(() => {
    setMessages([]);
    setSessionId(null);
    setError(null);
    setStepsExpanded({});
  }, [bibcode]);

  // 手动停止智能体执行
  const handleStop = async () => {
    if (!sessionId) return;
    try {
      await axios.post(`/api/chat/sessions/${sessionId}/stop`);
    } catch (e) {
      console.error('停止智能体执行失败:', e);
    }
  };

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

    setError(null);
    const userMsg: Message = {
      sender: 'user',
      text: questionText,
      imageUrl: pendingFigure?.url
    };
    setMessages(prev => [...prev, userMsg]);
    setInput('');
    setLoading(true);
    setShouldAutoScroll(true);

    const isFigureQuery = !!pendingFigure;
    const currentFigurePath = pendingFigure?.path;

    if (onClearPendingFigure) {
      onClearPendingFigure();
    }

    // AI 初始消息占位符
    const aiMsgPlaceholder: Message = {
      sender: 'ai',
      text: '',
      steps: [],
      sources: []
    };
    setMessages(prev => [...prev, aiMsgPlaceholder]);

    try {
      if (isFigureQuery && currentFigurePath) {
        // 图表多模态分析依然调用特定接口
        const res = await axios.post<{ answer: string }>('/api/chat/figure', {
          bibcode,
          image_path: currentFigurePath,
          question: questionText
        });
        setMessages(prev => {
          const next = [...prev];
          const last = next[next.length - 1];
          if (last && last.sender === 'ai') {
            last.text = res.data.answer;
          }
          return next;
        });
        setLoading(false);
      } else {
        // 调用智能体流式对话接口（带文献上下文提示，确保 Agent 优先分析当前阅读的文献）
        const contextPrefix = `[当前正在阅读文献: ${bibcode}] `;
        const requestQuery = questionText.includes(bibcode) ? questionText : `${contextPrefix}${questionText}`;

        try {
          const response = await fetch('/api/chat/agent', {
            method: 'POST',
            headers: {
              'Content-Type': 'application/json',
            },
            body: JSON.stringify({
              question: requestQuery,
              session_id: sessionId,
            }),
          });

          if (!response.ok) {
            throw new Error(`HTTP error! status: ${response.status}`);
          }

          const reader = response.body?.getReader();
          if (!reader) {
            throw new Error('ReadableStream not supported');
          }

          const decoder = new TextDecoder('utf-8');
          let buffer = '';

          while (true) {
            const { value, done } = await reader.read();
            if (done) break;

            buffer += decoder.decode(value, { stream: true });
            const lines = buffer.split('\n');
            buffer = lines.pop() || '';

            for (const line of lines) {
              const trimmed = line.trim();
              if (!trimmed) continue;

              if (trimmed.startsWith('data:')) {
                const dataStr = trimmed.slice(5).trim();
                if (!dataStr) continue;

                try {
                  const event = JSON.parse(dataStr);
                  
                  switch (event.type) {
                    case 'session':
                      if (!sessionId) {
                        setSessionId(event.session_id);
                      }
                      break;
                    case 'thought':
                      setMessages(prev => {
                        const next = [...prev];
                        const last = next[next.length - 1];
                        if (last && last.sender === 'ai') {
                          if (!last.steps) last.steps = [];
                          let stepObj = last.steps.find(s => s.step === event.step && s.type === 'thought');
                          if (!stepObj) {
                            stepObj = { step: event.step, type: 'thought', label: '思考推理中' };
                            last.steps.push(stepObj);
                          }
                          stepObj.detail = event.content;
                        }
                        return next;
                      });
                      break;
                    case 'tool_call':
                      setMessages(prev => {
                        const next = [...prev];
                        const last = next[next.length - 1];
                        if (last && last.sender === 'ai') {
                          if (!last.steps) last.steps = [];
                          let stepObj = last.steps.find(s => s.step === event.step && s.type === 'tool_call');
                          if (!stepObj) {
                            stepObj = { step: event.step, type: 'tool_call', label: `调用: ${getToolDisplayName(event.name)}`, isFinished: false };
                            last.steps.push(stepObj);
                          }
                          stepObj.detail = typeof event.arguments === 'string' ? event.arguments : JSON.stringify(event.arguments);
                        }
                        return next;
                      });
                      break;
                    case 'tool_result':
                      setMessages(prev => {
                        const next = [...prev];
                        const last = next[next.length - 1];
                        if (last && last.sender === 'ai') {
                          if (!last.steps) last.steps = [];
                          let stepObj = last.steps.find(s => s.step === event.step && s.type === 'tool_call');
                          if (stepObj) {
                            stepObj.label = `调用: ${getToolDisplayName(event.name)} (已获取观测数据)`;
                            stepObj.isError = event.is_error;
                            stepObj.isFinished = true;
                          }
                          
                          // 提取 RAG 检索到的参考来源文献，并在 sources 中进行去重和追加
                          if (event.name === 'rag_search' && event.metadata && event.metadata.sources) {
                            const newSources: RetrievalSource[] = event.metadata.sources.map((s: any) => ({
                              bibcode: s.bibcode,
                              paragraph_index: s.paragraph_index,
                              content: s.preview || '',
                              distance: s.distance,
                            }));
                            
                            const existing = last.sources || [];
                            const combined = [...existing];
                            
                            for (const src of newSources) {
                              if (!combined.some(c => c.bibcode === src.bibcode && c.paragraph_index === src.paragraph_index)) {
                                combined.push(src);
                              }
                            }
                            last.sources = combined;
                          }
                        }
                        return next;
                      });
                      break;
                    case 'text_delta':
                      setMessages(prev => {
                        const next = [...prev];
                        const last = next[next.length - 1];
                        if (last && last.sender === 'ai') {
                          last.text += event.content;
                        }
                        return next;
                      });
                      break;
                    case 'error':
                      setMessages(prev => {
                        const next = [...prev];
                        const last = next[next.length - 1];
                        if (last && last.sender === 'ai') {
                          if (!last.steps) last.steps = [];
                          last.steps.push({
                            step: 99,
                            type: 'error',
                            label: `错误返回`,
                            detail: event.message
                          });
                        }
                        return next;
                      });
                      break;
                  }
                } catch (e) {
                  console.error('解析流式 SSE 失败:', e);
                }
              }
            }
          }
          setLoading(false);

        } catch (streamErr) {
          console.warn('智能体流式对话失败，正在尝试回退至旧版本 RAG 单次问答...', streamErr);
          
          // 降级回退 (Deprecated Fallback) —— 调用原始单次 /api/chat/rag
          const res = await axios.post<{ answer: string; sources: RetrievalSource[] }>('/api/chat/rag', {
            question: questionText,
            top_k: 5
          });
          
          setMessages(prev => {
            const next = [...prev];
            const last = next[next.length - 1];
            if (last && last.sender === 'ai') {
              last.text = res.data.answer;
              last.sources = res.data.sources;
            }
            return next;
          });
          setLoading(false);
        }
      }
    } catch (err: any) {
      console.error('问答请求失败:', err);
      setError(err.response?.data?.error || err.message || '网络请求错误，请稍后重试');
      // 如果报错，清空多余的消息占位符
      setMessages(prev => {
        const next = [...prev];
        if (next.length > 0 && next[next.length - 1].sender === 'ai' && !next[next.length - 1].text) {
          next.pop();
        }
        return next;
      });
      setLoading(false);
    }
  };

  const toggleSteps = (msgIdx: number) => {
    setStepsExpanded(prev => ({
      ...prev,
      [msgIdx]: !prev[msgIdx]
    }));
  };

  return (
    <div className="console-panel rounded-xl border border-slate-200 bg-slate-50 flex flex-col overflow-hidden relative shadow-sm h-full">
      {/* 头部面板 */}
      <div className="px-4 py-3.5 border-b border-slate-200 flex items-center justify-between bg-white shrink-0">
        <div className="flex items-center gap-1.5 select-none">
          <Sparkles className="w-4 h-4 text-sky-600 animate-pulse" />
          <span className="text-xs font-bold text-slate-800">科研智能问答助手</span>
        </div>
        <button onClick={onClose} className="text-slate-400 hover:text-slate-600 transition-colors cursor-pointer">
          <X className="w-4 h-4" />
        </button>
      </div>

      {/* 对话消息区 */}
      <div 
        ref={scrollContainerRef}
        onScroll={handleScroll}
        className="flex-1 overflow-y-auto p-4 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">
                基于 ReAct 框架的智能科研体。支持自动读取已解析的 Markdown 文献，在向量库进行 RAG 查询以及天体参数检索。
              </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-100 border border-slate-200 rounded-lg p-3 text-xs text-slate-700 leading-relaxed font-medium transition-all shadow-sm cursor-pointer hover:border-sky-300"
                  >
                    {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-xl px-4 py-3 text-xs leading-relaxed font-medium shadow-sm border ${
                  msg.sender === 'user'
                    ? 'bg-sky-600 text-white border-sky-600 select-text'
                    : 'bg-white text-slate-800 border-slate-200 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">
                    {/* Collapsible Steps list */}
                    {msg.steps && msg.steps.length > 0 && (
                      <div className="border border-slate-100 bg-slate-50/50 rounded-lg p-2 select-none">
                        <button
                          onClick={() => toggleSteps(index)}
                          className="flex items-center justify-between w-full text-left text-[10px] font-bold text-slate-500 hover:text-slate-700 transition-colors cursor-pointer"
                        >
                          <span className="flex items-center gap-1">
                            <Brain className="w-3.5 h-3.5 text-purple-500" />
                            <span>推理与工具链展开 ({msg.steps.length} 步)</span>
                          </span>
                          {stepsExpanded[index] ? <ChevronUp className="w-3.5 h-3.5" /> : <ChevronDown className="w-3.5 h-3.5" />}
                        </button>
                        
                        {stepsExpanded[index] && (
                          <div className="mt-2 pl-2.5 border-l border-slate-200 space-y-1.5 transition-all">
                            {msg.steps.map((step, sIdx) => (
                              <div key={sIdx} className="text-[9px] text-slate-655 flex flex-col gap-0.5">
                                <span className="font-bold flex items-center gap-1">
                                  {step.type === 'thought' ? (
                                    <Brain className="w-3 h-3 text-purple-500 shrink-0" />
                                  ) : step.type === 'error' ? (
                                    <AlertTriangle className="w-3 h-3 text-rose-500 shrink-0" />
                                  ) : (
                                    <Settings className={`w-3 h-3 text-sky-600 shrink-0 ${step.isFinished ? '' : 'animate-spin'}`} />
                                  )}
                                  <span>{step.label}</span>
                                </span>
                                {step.detail && (
                                  <span className="font-mono text-[9px] opacity-75 whitespace-pre-wrap block bg-white px-1.5 py-0.5 rounded border border-slate-100 max-h-24 overflow-y-auto">
                                    {step.detail}
                                  </span>
                                )}
                              </div>
                            ))}
                          </div>
                        )}
                      </div>
                    )}

                    {/* Final output text */}
                    {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">
                        <ReactMarkdown
                          remarkPlugins={[remarkMath, remarkGfm]}
                          rehypePlugins={[rehypeRaw, [rehypeSanitize, safeSchema], rehypeKatex]}
                        >
                          {msg.text}
                        </ReactMarkdown>
                      </div>
                    ) : (
                      !loading && <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>
          ))
        )}

        {loading && 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>

      {/* 待提问图片预览 */}
      {pendingFigure && (
        <div className="px-3 py-2 bg-slate-100 border-t border-slate-200 flex items-center justify-between shrink-0 select-none">
          <div className="flex items-center gap-2">
            <img
              src={pendingFigure.url}
              alt="Preview"
              className="w-10 h-10 object-cover rounded-md border border-slate-300"
            />
            <div className="text-[10px] text-slate-550">
              <span className="font-bold text-slate-700 block">已选中图表插图</span>
              <span className="font-mono text-slate-400 truncate max-w-48 block">{pendingFigure.path.split('/').pop()}</span>
            </div>
          </div>
          <button
            type="button"
            onClick={onClearPendingFigure}
            className="text-slate-400 hover:text-slate-655 p-1 cursor-pointer"
          >
            <X className="w-3.5 h-3.5" />
          </button>
        </div>
      )}

      {/* 底栏输入区 */}
      <form
        onSubmit={(e) => {
          e.preventDefault();
          handleSend(input);
        }}
        className="p-3 border-t border-slate-200 bg-white shrink-0"
      >
        <div className="flex gap-2 relative items-center">
          <input
            type="text"
            value={input}
            onChange={(e) => setInput(e.target.value)}
            disabled={loading}
            placeholder={pendingFigure ? "针对选中图表提问..." : "向 AI 馆藏助手提问..."}
            className="flex-1 bg-slate-50 border border-slate-200 rounded-xl text-xs text-slate-900 placeholder-slate-400 pl-3 pr-10 py-2.5 focus:outline-none focus:bg-white focus:border-sky-500 focus:ring-1 focus:ring-sky-500/10 leading-relaxed font-semibold transition-all disabled:opacity-60"
          />
          {loading ? (
            <button
              type="button"
              onClick={handleStop}
              className="absolute right-1.5 p-1.5 rounded-lg bg-rose-600 hover:bg-rose-700 text-white transition-colors cursor-pointer flex items-center justify-center"
              title="手动停止执行"
            >
              <Square className="w-3.5 h-3.5 fill-white" />
            </button>
          ) : (
            <button
              type="submit"
              disabled={!input.trim()}
              className="absolute right-1.5 p-1.5 rounded-lg bg-sky-600 hover:bg-sky-700 text-white disabled:bg-slate-200 disabled:text-slate-400 transition-colors cursor-pointer flex items-center justify-center"
            >
              <Send className="w-3.5 h-3.5" />
            </button>
          )}
        </div>
      </form>
    </div>
  );
}