import { useCallback, useEffect, useRef, useState } from 'react';
import {
  Zap,
  X,
  Plus,
  History,
  MessageSquare,
  Send,
  Trash2,
} from 'lucide-react';
import { useChatStore } from '../../lib/chat-store';
import type { ChatMessage } from '../../lib/chat-store';
import { streamChat } from '../../lib/chat-api';
import { useNotebookStore } from '../../lib/notebook-store';
import { marked } from 'marked';
import markedKatex from 'marked-katex-extension';
import DOMPurify from 'dompurify';
import 'katex/dist/katex.min.css';
import '../../styles/chat-widget.css';

// ── Markdown rendering (marked + KaTeX + DOMPurify) ──────

marked.setOptions({
  breaks: true,
  gfm: true,
});

marked.use(markedKatex({
  throwOnError: false,
  nonStandard: true,
}));

// KaTeX generates SVG + spans with specific classes/attributes.
// DOMPurify must allow these through.
const KATEX_TAGS = [
  'math', 'semantics', 'mrow', 'mi', 'mo', 'mn', 'msup', 'msub',
  'mfrac', 'munderover', 'mover', 'munder', 'msqrt', 'mroot',
  'mtable', 'mtr', 'mtd', 'mspace', 'mtext', 'menclose',
  'annotation', 'annotation-xml',
];

const KATEX_ATTRS = [
  'xmlns', 'encoding', 'mathvariant', 'displaystyle', 'scriptlevel',
  'fence', 'stretchy', 'symmetric', 'lspace', 'rspace',
  'linethickness', 'columnalign', 'rowalign', 'columnspacing',
  'rowspacing', 'columnlines', 'rowlines', 'frame', 'framespacing',
  'width', 'height', 'voffset', 'accent', 'accentunder',
  'notation', 'minsize', 'maxsize', 'movablelimits',
  'aria-hidden', 'focusable', 'role', 'tabindex',
  'viewBox', 'preserveAspectRatio', 'd', 'fill', 'stroke',
  'stroke-width', 'stroke-linecap', 'stroke-linejoin',
  'transform', 'clip-path',
];

// LLMs often emit $$...$$ inline (e.g. "is:$$\frac{1}{s}$$\nNext")
// but marked-katex-extension requires $$ on its own line for display mode.
// This normalizer ensures $$ delimiters get their own lines.
function normalizeDisplayMath(text: string): string {
  return text.replace(/\$\$([\s\S]*?)\$\$/g, (_match, inner: string) => {
    const trimmed = inner.trim();
    return `\n$$\n${trimmed}\n$$\n`;
  });
}

function renderMarkdown(text: string): string {
  const normalized = normalizeDisplayMath(text);
  const raw = marked.parse(normalized, { async: false }) as string;
  return DOMPurify.sanitize(raw, {
    ADD_TAGS: KATEX_TAGS,
    ADD_ATTR: [...KATEX_ATTRS, 'target', 'class', 'style'],
  });
}

// ── Components ──────────────────────────────────────────

function MessageBubble({ msg, isStreaming }: { msg: ChatMessage; isStreaming?: boolean }) {
  return (
    <div className={`chat-bubble ${msg.role}`}>
      {msg.role === 'assistant' ? (
        <>
          <span dangerouslySetInnerHTML={{ __html: renderMarkdown(msg.text) }} />
          {isStreaming && <span className="chat-cursor" />}
        </>
      ) : (
        msg.text
      )}
    </div>
  );
}

function HistoryView({
  onBack,
}: {
  onBack: () => void;
}) {
  const conversations = useChatStore((s) => s.conversations);
  const activeId = useChatStore((s) => s.activeConversationId);
  const setActive = useChatStore((s) => s.setActiveConversation);
  const deleteConv = useChatStore((s) => s.deleteConversation);
  const [pendingDeleteId, setPendingDeleteId] = useState<string | null>(null);
  const pendingTimer = useRef<ReturnType<typeof setTimeout> | null>(null);

  const handleDelete = useCallback(
    (id: string, e: React.MouseEvent) => {
      e.stopPropagation();
      if (pendingDeleteId === id) {
        // Second click — confirm
        if (pendingTimer.current) clearTimeout(pendingTimer.current);
        setPendingDeleteId(null);
        deleteConv(id);
      } else {
        // First click — arm
        if (pendingTimer.current) clearTimeout(pendingTimer.current);
        setPendingDeleteId(id);
        pendingTimer.current = setTimeout(() => setPendingDeleteId(null), 3000);
      }
    },
    [pendingDeleteId, deleteConv],
  );

  useEffect(() => {
    return () => {
      if (pendingTimer.current) clearTimeout(pendingTimer.current);
    };
  }, []);

  if (conversations.length === 0) {
    return (
      <div className="chat-history-list">
        <div className="chat-empty">No conversations yet</div>
      </div>
    );
  }

  return (
    <div className="chat-history-list">
      {conversations.map((c) => (
        <button
          key={c.id}
          className={`chat-history-item ${c.id === activeId ? 'active' : ''}`}
          onClick={() => {
            setActive(c.id);
            onBack();
          }}
        >
          <MessageSquare size={14} />
          <span className="chat-history-title">{c.title}</span>
          <span className="chat-history-count">{c.messages.length}</span>
          <span
            className={`chat-history-delete ${pendingDeleteId === c.id ? 'confirm' : ''}`}
            onClick={(e) => handleDelete(c.id, e)}
            title={pendingDeleteId === c.id ? 'Click again to delete' : 'Delete conversation'}
          >
            <Trash2 size={12} />
          </span>
        </button>
      ))}
    </div>
  );
}

// ── Main Widget ─────────────────────────────────────────

export default function ChatWidget() {
  const panelOpen = useChatStore((s) => s.panelOpen);
  const streaming = useChatStore((s) => s.streaming);
  const togglePanel = useChatStore((s) => s.togglePanel);
  const closePanel = useChatStore((s) => s.closePanel);
  const createConversation = useChatStore((s) => s.createConversation);
  const addUserMessage = useChatStore((s) => s.addUserMessage);
  const addAssistantMessage = useChatStore((s) => s.addAssistantMessage);
  const appendToLastAssistant = useChatStore((s) => s.appendToLastAssistant);
  const setStreaming = useChatStore((s) => s.setStreaming);
  const getActiveConversation = useChatStore((s) => s.getActiveConversation);
  const activeConversationId = useChatStore((s) => s.activeConversationId);

  // Notebook context from the notebook editor (if on a notebook page)
  const notebookId = useNotebookStore((s) => s.notebookId);
  const notebook = useNotebookStore((s) => s.notebook);

  const [input, setInput] = useState('');
  const [statusText, setStatusText] = useState('');
  const [showHistory, setShowHistory] = useState(false);
  const [reasoningText, setReasoningText] = useState('');
  const [reasoningTime, setReasoningTime] = useState(0);

  const messagesEndRef = useRef<HTMLDivElement>(null);
  const inputRef = useRef<HTMLInputElement>(null);
  const abortRef = useRef<AbortController | null>(null);

  // Token batching: accumulate tokens in a ref, flush to state once per animation frame.
  // Without this, React 19 batches all rapid set() calls from the SSE loop into one
  // render at stream end — so the user sees nothing until the LLM finishes.
  const pendingTokensRef = useRef('');
  const flushRafRef = useRef(0);

  const activeConv = getActiveConversation();
  const messages = activeConv?.messages ?? [];

  // Auto-scroll on new messages — use direct scrollTop instead of scrollIntoView
  // because scrollIntoView walks up the DOM and scrolls ALL ancestors (including
  // the panel with overflow:hidden), which pushes the header off-screen.
  useEffect(() => {
    const container = messagesEndRef.current?.parentElement;
    if (container) {
      container.scrollTo({ top: container.scrollHeight, behavior: 'smooth' });
    }
  }, [messages.length, streaming]);

  // Focus input when panel opens
  useEffect(() => {
    if (panelOpen && !showHistory) {
      setTimeout(() => inputRef.current?.focus(), 100);
    }
  }, [panelOpen, showHistory]);

  // Keyboard: Escape to close
  useEffect(() => {
    function handleEscape(e: KeyboardEvent) {
      if (e.key === 'Escape' && panelOpen) {
        closePanel();
      }
    }
    document.addEventListener('keydown', handleEscape);
    return () => document.removeEventListener('keydown', handleEscape);
  }, [panelOpen, closePanel]);

  const handleSend = useCallback(async () => {
    const question = input.trim();
    if (!question || streaming) return;

    // Ensure there's an active conversation
    if (!activeConversationId) {
      createConversation();
    }

    setInput('');
    addUserMessage(question);
    addAssistantMessage('');
    setStreaming(true);
    setStatusText('');
    setReasoningText('');
    setReasoningTime(0);

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

    const reasoningStart = Date.now();

    try {
      const notebookCtx = notebookId && notebook
        ? {
            notebook_id: notebookId,
            title: notebook.metadata.title,
            engine: notebook.metadata.engine,
          }
        : null;

      for await (const evt of streamChat({
        question,
        notebook: notebookCtx,
        signal: abortController.signal,
      })) {
        switch (evt.event) {
          case 'status':
            setStatusText((evt.data as { text: string }).text);
            break;
          case 'token':
            // Accumulate in ref; flush to Zustand once per animation frame
            pendingTokensRef.current += (evt.data as { text: string }).text;
            if (!flushRafRef.current) {
              flushRafRef.current = requestAnimationFrame(() => {
                if (pendingTokensRef.current) {
                  appendToLastAssistant(pendingTokensRef.current);
                  pendingTokensRef.current = '';
                }
                flushRafRef.current = 0;
              });
              setStatusText('');
            }
            break;
          case 'reasoning':
            setReasoningText((prev) => prev + (evt.data as { text: string }).text);
            setReasoningTime(Math.round((Date.now() - reasoningStart) / 1000));
            break;
          case 'error':
            appendToLastAssistant(
              `\n\n*Error: ${(evt.data as { text: string }).text}*`,
            );
            break;
          case 'done':
            break;
        }
      }
    } catch (err) {
      if (err instanceof DOMException && err.name === 'AbortError') {
        // User cancelled — that's fine
      } else {
        appendToLastAssistant(
          `\n\n*Error: ${err instanceof Error ? err.message : 'Connection failed'}*`,
        );
      }
    } finally {
      // Flush any remaining buffered tokens
      if (flushRafRef.current) {
        cancelAnimationFrame(flushRafRef.current);
        flushRafRef.current = 0;
      }
      if (pendingTokensRef.current) {
        appendToLastAssistant(pendingTokensRef.current);
        pendingTokensRef.current = '';
      }
      setStreaming(false);
      setStatusText('');
      abortRef.current = null;
    }
  }, [
    input, streaming, activeConversationId, notebookId, notebook,
    createConversation, addUserMessage, addAssistantMessage,
    appendToLastAssistant, setStreaming,
  ]);

  const handleKeyDown = useCallback(
    (e: React.KeyboardEvent) => {
      if (e.key === 'Enter' && !e.shiftKey) {
        e.preventDefault();
        handleSend();
      }
    },
    [handleSend],
  );

  const handleNewConversation = useCallback(() => {
    // If streaming, abort first
    if (abortRef.current) {
      abortRef.current.abort();
      abortRef.current = null;
      setStreaming(false);
    }
    createConversation();
    setShowHistory(false);
    setReasoningText('');
    setReasoningTime(0);
    setTimeout(() => inputRef.current?.focus(), 100);
  }, [createConversation, setStreaming]);

  return (
    <>
      {/* Floating action button */}
      <button
        className={`chat-fab ${panelOpen ? 'active' : ''}`}
        onClick={togglePanel}
        aria-label="Toggle chat"
        title="Circuit assistant"
      >
        {panelOpen ? <X size={20} /> : <Zap size={20} />}
      </button>

      {/* Chat panel */}
      {panelOpen && (
        <div className="chat-panel">
          {/* Header */}
          <div className="chat-header">
            <Zap size={16} className="text-blue-500" />
            <span className="chat-header-title">
              {showHistory
                ? 'Conversations'
                : activeConv?.title ?? 'Circuit Assistant'}
            </span>
            <button
              className="chat-header-btn"
              onClick={handleNewConversation}
              title="New conversation"
            >
              <Plus size={16} />
            </button>
            <button
              className="chat-header-btn"
              onClick={() => setShowHistory(!showHistory)}
              title={showHistory ? 'Back to chat' : 'Conversation history'}
            >
              {showHistory ? <MessageSquare size={16} /> : <History size={16} />}
            </button>
            <button
              className="chat-header-btn"
              onClick={closePanel}
              title="Close"
            >
              <X size={16} />
            </button>
          </div>

          {/* Body */}
          {showHistory ? (
            <HistoryView onBack={() => setShowHistory(false)} />
          ) : (
            <>
              <div className="chat-messages">
                {messages.length === 0 && !streaming && (
                  <div className="chat-empty">
                    Ask about circuits, netlists, or simulation results
                  </div>
                )}
                {messages.map((msg, i) => (
                  <MessageBubble
                    key={`${msg.timestamp}-${i}`}
                    msg={msg}
                    isStreaming={
                      streaming &&
                      i === messages.length - 1 &&
                      msg.role === 'assistant'
                    }
                  />
                ))}
                {reasoningText && streaming && (
                  <details className="chat-reasoning">
                    <summary>Thinking{reasoningTime > 0 ? ` (${reasoningTime}s)` : '…'}</summary>
                    <div className="chat-reasoning-body">{reasoningText}</div>
                  </details>
                )}
                {statusText && (
                  <div className="chat-status">{statusText}</div>
                )}
                <div ref={messagesEndRef} />
              </div>

              {/* Input */}
              <div className="chat-input-bar">
                <input
                  ref={inputRef}
                  className="chat-input"
                  type="text"
                  placeholder={
                    notebookId
                      ? 'Ask about this circuit…'
                      : 'Ask a circuit question…'
                  }
                  value={input}
                  onChange={(e) => setInput(e.target.value)}
                  onKeyDown={handleKeyDown}
                  disabled={streaming}
                />
                <button
                  className="chat-send-btn"
                  onClick={handleSend}
                  disabled={streaming || !input.trim()}
                  title="Send"
                >
                  <Send size={16} />
                </button>
              </div>
            </>
          )}
        </div>
      )}
    </>
  );
}