"use client"

import { useState, useRef, useEffect, useCallback, memo, Fragment } from "react"
import { MessageSquare, Send, X, Loader2, User, Bot, Copy, Check, ChevronRight } from "lucide-react"
import ReactMarkdown from "react-markdown"
import remarkGfm from "remark-gfm"
import { CodeHighlighter, CODE_STYLE, CODE_STYLE_SMALL } from "./code-highlighter"

interface ChatMessage {
  role: "user" | "assistant"
  content: string
}

interface ToolStep {
  tool: string
  displayName: string
  status: "running" | "done"
  summary?: string
  content?: string
}

interface TimelineChatProps {
  traceId: string
  onClose: () => void
}

export const TimelineChat = memo(function TimelineChat({
  traceId,
  onClose,
}: TimelineChatProps) {
  const [messages, setMessages] = useState<ChatMessage[]>([])
  const [input, setInput] = useState("")
  const [isStreaming, setIsStreaming] = useState(false)
  const [completedRounds, setCompletedRounds] = useState<ToolStep[][]>([])
  const [activeSteps, setActiveSteps] = useState<ToolStep[]>([])
  const [statusMessage, setStatusMessage] = useState<string | null>(null)
  const messagesEndRef = useRef<HTMLDivElement>(null)
  const inputRef = useRef<HTMLTextAreaElement>(null)
  const abortRef = useRef<AbortController | null>(null)

  const scrollToBottom = useCallback(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" })
  }, [])

  useEffect(() => {
    scrollToBottom()
  }, [messages, completedRounds, activeSteps, statusMessage, scrollToBottom])

  useEffect(() => {
    requestAnimationFrame(() => inputRef.current?.focus())
  }, [])

  const sendMessage = useCallback(async () => {
    const trimmed = input.trim()
    if (!trimmed || isStreaming) return

    const userMessage: ChatMessage = { role: "user", content: trimmed }
    const newMessages = [...messages, userMessage]
    setMessages(newMessages)
    setInput("")
    setIsStreaming(true)

    const controller = new AbortController()
    abortRef.current = controller

    setMessages((prev) => [...prev, { role: "assistant", content: "" }])
    setCompletedRounds([])
    setActiveSteps([])

    try {
      const res = await fetch("/api/observability/chat", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ traceId, messages: newMessages }),
        signal: controller.signal,
      })

      if (!res.ok) {
        const errData = await res.json().catch(() => ({}))
        throw new Error(errData.error || `Request failed: ${res.status}`)
      }

      const reader = res.body?.getReader()
      if (!reader) throw new Error("No response stream")

      const decoder = new TextDecoder()
      let buffer = ""

      while (true) {
        const { done, value } = 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) {
          if (!line.startsWith("data: ")) continue
          const data = line.slice(6).trim()
          if (data === "[DONE]") continue

          try {
            const parsed = JSON.parse(data)

            if (parsed.type === "status") {
              setStatusMessage(parsed.message ?? null)
              continue
            }

            // Handle typed events (new protocol)
            if (parsed.type === "tool_start") {
              setStatusMessage(null)
              setActiveSteps((prev) => {
                // If all previous steps are done, this is a new round — commit previous steps
                if (prev.length > 0 && prev.every((s) => s.status === "done")) {
                  setCompletedRounds((rounds) => [...rounds, prev])
                  return [{
                    tool: parsed.tool,
                    displayName: parsed.displayName ?? parsed.tool,
                    status: "running",
                  }]
                }
                return [
                  ...prev,
                  {
                    tool: parsed.tool,
                    displayName: parsed.displayName ?? parsed.tool,
                    status: "running",
                  },
                ]
              })
              continue
            }

            if (parsed.type === "tool_result") {
              setActiveSteps((prev) =>
                prev.map((step) =>
                  step.tool === parsed.tool && step.status === "running"
                    ? { ...step, status: "done", summary: parsed.summary, content: parsed.content }
                    : step,
                ),
              )
              continue
            }

            // Handle text — both new {type: "text", text} and old {text} formats
            const textContent = parsed.type === "text" ? parsed.text : parsed.text
            if (textContent) {
              setStatusMessage(null)
              setMessages((prev) => {
                const updated = [...prev]
                const last = updated[updated.length - 1]
                if (last?.role === "assistant") {
                  updated[updated.length - 1] = {
                    ...last,
                    content: last.content + textContent,
                  }
                }
                return updated
              })
            }
            if (parsed.error) {
              throw new Error(parsed.error)
            }
          } catch (e) {
            if (e instanceof SyntaxError) continue
            throw e
          }
        }
      }
    } catch (err) {
      if ((err as Error).name === "AbortError") return

      setMessages((prev) => {
        const updated = [...prev]
        const last = updated[updated.length - 1]
        if (last?.role === "assistant" && !last.content) {
          updated[updated.length - 1] = {
            ...last,
            content: `Error: ${(err as Error).message}`,
          }
        }
        return updated
      })
    } finally {
      // Commit any remaining active steps as a final completed round
      setActiveSteps((prev) => {
        if (prev.length > 0) {
          setCompletedRounds((rounds) => [...rounds, prev])
        }
        return []
      })
      setStatusMessage(null)
      setIsStreaming(false)
      abortRef.current = null
    }
  }, [input, isStreaming, messages, traceId])

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

  const stopStreaming = useCallback(() => {
    abortRef.current?.abort()
  }, [])

  const [copied, setCopied] = useState(false)
  const exportChat = useCallback(() => {
    const text = messages
      .filter((m) => m.content)
      .map((m) => `${m.role === "user" ? "User" : "Assistant"}: ${m.content}`)
      .join("\n\n")
    navigator.clipboard.writeText(text).then(() => {
      setCopied(true)
      setTimeout(() => setCopied(false), 2000)
    })
  }, [messages])

  return (
    <div className="flex flex-col h-full bg-white dark:bg-zinc-900 border-l border-zinc-200 dark:border-zinc-700">
      <div className="flex items-center justify-between px-4 py-3 border-b border-zinc-200 dark:border-zinc-700 bg-zinc-50 dark:bg-zinc-800/50">
        <div className="flex items-center gap-2">
          <MessageSquare className="h-4 w-4 text-zinc-500" />
          <span className="text-sm font-medium text-zinc-900 dark:text-white">
            Chat with Trace
          </span>
        </div>
        <div className="flex items-center gap-1">
          {messages.length > 0 && (
            <button
              onClick={exportChat}
              className="p-1.5 rounded text-zinc-400 hover:text-zinc-600 dark:hover:text-zinc-300 hover:bg-zinc-200 dark:hover:bg-zinc-700 transition-colors"
              title="Copy chat to clipboard"
            >
              {copied ? <Check className="h-4 w-4 text-green-500" /> : <Copy className="h-4 w-4" />}
            </button>
          )}
          <button
            onClick={onClose}
            className="p-1.5 rounded text-zinc-400 hover:text-zinc-600 dark:hover:text-zinc-300 hover:bg-zinc-200 dark:hover:bg-zinc-700 transition-colors"
          >
            <X className="h-4 w-4" />
          </button>
        </div>
      </div>

      <div className="flex-1 overflow-y-auto px-4 py-4 space-y-4">
        {messages.length === 0 && (
          <div className="flex flex-col items-center justify-center h-full text-center px-6">
            <div className="w-12 h-12 rounded-full bg-zinc-100 dark:bg-zinc-800 flex items-center justify-center mb-3">
              <MessageSquare className="h-6 w-6 text-zinc-400" />
            </div>
            <p className="text-sm text-zinc-500 dark:text-zinc-400 mb-1">
              Ask about this optimization trace
            </p>
            <p className="text-xs text-zinc-400 dark:text-zinc-500">
              e.g. &quot;Why was candidate 1 ranked best?&quot; or &quot;What optimizations were attempted?&quot;
            </p>
          </div>
        )}

        {messages.map((msg, i) => {
          const isLastAssistant = i === messages.length - 1 && msg.role === "assistant"

          if (isLastAssistant) {
            return (
              <Fragment key={i}>
                {completedRounds.map((steps, j) => (
                  <ToolRoundBubble key={`round-${j}`} steps={steps} />
                ))}
                {activeSteps.length > 0 && <ToolRoundBubble steps={activeSteps} isActive />}
                {statusMessage && (
                  <div className="flex items-center gap-2 text-xs text-zinc-400 dark:text-zinc-500 px-2">
                    <Loader2 className="h-3 w-3 animate-spin" />
                    <span>{statusMessage}</span>
                  </div>
                )}
                <ChatBubble message={msg} isStreaming={isStreaming} />
              </Fragment>
            )
          }

          return <ChatBubble key={i} message={msg} isStreaming={false} />
        })}

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

      <div className="border-t border-zinc-200 dark:border-zinc-700 px-4 py-3">
        <div className="flex items-end gap-2">
          <textarea
            ref={inputRef}
            value={input}
            onChange={(e) => setInput(e.target.value)}
            onKeyDown={handleKeyDown}
            placeholder="Ask about this trace..."
            rows={1}
            className="flex-1 resize-none rounded-lg border border-zinc-200 dark:border-zinc-700 bg-white dark:bg-zinc-800 px-3 py-2 text-sm text-zinc-900 dark:text-zinc-100 placeholder-zinc-400 focus:outline-none focus:ring-2 focus:ring-blue-500 dark:focus:ring-blue-400 focus:border-transparent max-h-32 min-h-[38px]"
            style={{ height: "38px" }}
            onInput={(e) => {
              const target = e.target as HTMLTextAreaElement
              target.style.height = "38px"
              target.style.height = `${Math.min(target.scrollHeight, 128)}px`
            }}
            disabled={isStreaming}
          />
          {isStreaming ? (
            <button
              onClick={stopStreaming}
              className="flex-shrink-0 p-2 rounded-lg bg-red-500 hover:bg-red-600 text-white transition-colors"
              title="Stop generating"
            >
              <X className="h-4 w-4" />
            </button>
          ) : (
            <button
              onClick={sendMessage}
              disabled={!input.trim()}
              className="flex-shrink-0 p-2 rounded-lg bg-blue-600 hover:bg-blue-700 disabled:bg-zinc-300 dark:disabled:bg-zinc-700 text-white disabled:text-zinc-500 transition-colors"
              title="Send message"
            >
              <Send className="h-4 w-4" />
            </button>
          )}
        </div>
      </div>
    </div>
  )
})

const ToolRoundBubble = memo(function ToolRoundBubble({
  steps,
  isActive = false,
}: {
  steps: ToolStep[]
  isActive?: boolean
}) {
  const allDone = steps.every((s) => s.status === "done")
  const [expandedSteps, setExpandedSteps] = useState<Set<number>>(new Set())

  const toggleStep = useCallback((index: number) => {
    setExpandedSteps((prev) => {
      const next = new Set(prev)
      if (next.has(index)) {
        next.delete(index)
      } else {
        next.add(index)
      }
      return next
    })
  }, [])

  return (
    <div className="flex gap-3">
      <div
        className={`flex-shrink-0 w-7 h-7 rounded-full flex items-center justify-center ${
          isActive && !allDone
            ? "bg-amber-100 dark:bg-amber-900/40"
            : "bg-zinc-100 dark:bg-zinc-800"
        }`}
      >
        {isActive && !allDone ? (
          <Loader2 className="h-3.5 w-3.5 animate-spin text-amber-600 dark:text-amber-400" />
        ) : (
          <Check className="h-3.5 w-3.5 text-green-500" />
        )}
      </div>
      <div className="max-w-[85%] rounded-lg px-3 py-2 bg-zinc-50 dark:bg-zinc-800/50 border border-zinc-200 dark:border-zinc-700">
        <div className="flex flex-col gap-1">
          {steps.map((step, i) => {
            const isExpanded = expandedSteps.has(i)
            const hasContent = step.status === "done" && !!step.content

            return (
              <div key={`${step.tool}-${i}`}>
                <div className="flex items-center gap-2 text-xs text-zinc-500 dark:text-zinc-400">
                  {step.status === "running" ? (
                    <Loader2 className="h-3 w-3 animate-spin flex-shrink-0 text-amber-500" />
                  ) : (
                    <Check className="h-3 w-3 text-green-500 flex-shrink-0" />
                  )}
                  {hasContent ? (
                    <button
                      onClick={() => toggleStep(i)}
                      className="flex items-center gap-1 hover:text-zinc-700 dark:hover:text-zinc-300 transition-colors text-left"
                      aria-expanded={isExpanded}
                    >
                      <ChevronRight
                        className={`h-3 w-3 flex-shrink-0 transition-transform ${isExpanded ? "rotate-90" : ""}`}
                      />
                      <span>
                        {step.displayName}
                        {step.summary && (
                          <span className="text-zinc-400 dark:text-zinc-500"> — {step.summary}</span>
                        )}
                      </span>
                    </button>
                  ) : (
                    <span>
                      {step.displayName}
                      {step.summary && (
                        <span className="text-zinc-400 dark:text-zinc-500"> — {step.summary}</span>
                      )}
                    </span>
                  )}
                </div>
                {isExpanded && step.content && (
                  <div className="mt-1 ml-5">
                    <ToolResultContent content={step.content} />
                  </div>
                )}
              </div>
            )
          })}
        </div>
      </div>
    </div>
  )
})

const ToolResultContent = memo(function ToolResultContent({ content }: { content: string }) {
  return (
    <div className="max-h-[300px] overflow-y-auto rounded border border-zinc-200 dark:border-zinc-700 bg-zinc-100 dark:bg-zinc-800">
      <div className="prose prose-sm dark:prose-invert max-w-none prose-p:my-1 prose-pre:my-1 prose-headings:my-1 text-xs">
        <ReactMarkdown
          remarkPlugins={[remarkGfm]}
          components={{
            code({ className, children, ...props }) {
              const match = /language-(\w+)/.exec(className || "")
              const codeStr = String(children).replace(/\n$/, "")

              if (match) {
                return (
                  <div className="rounded overflow-hidden my-1 not-prose">
                    <CodeHighlighter
                      language={match[1]}
                      code={codeStr}
                      customStyle={CODE_STYLE_SMALL}
                    />
                  </div>
                )
              }

              return (
                <code
                  className="px-1 py-0.5 rounded bg-zinc-200 dark:bg-zinc-700 text-xs font-mono"
                  {...props}
                >
                  {children}
                </code>
              )
            },
            // Render plain text content as preformatted when no markdown structure detected
            p({ children }) {
              return <p className="whitespace-pre-wrap break-words text-xs text-zinc-600 dark:text-zinc-300">{children}</p>
            },
          }}
        >
          {content}
        </ReactMarkdown>
      </div>
    </div>
  )
})

const ChatBubble = memo(function ChatBubble({
  message,
  isStreaming,
}: {
  message: ChatMessage
  isStreaming: boolean
}) {
  const isUser = message.role === "user"

  return (
    <div className={`flex gap-3 ${isUser ? "justify-end" : ""}`}>
      {!isUser && (
        <div className="flex-shrink-0 w-7 h-7 rounded-full bg-blue-100 dark:bg-blue-900/40 flex items-center justify-center">
          <Bot className="h-3.5 w-3.5 text-blue-600 dark:text-blue-400" />
        </div>
      )}
      <div
        className={`max-w-[85%] rounded-lg px-3 py-2 text-sm ${
          isUser
            ? "bg-blue-600 text-white"
            : "bg-zinc-100 dark:bg-zinc-800 text-zinc-900 dark:text-zinc-100"
        }`}
      >
        {isUser ? (
          <p className="whitespace-pre-wrap">{message.content}</p>
        ) : message.content ? (
          <div className="prose prose-sm dark:prose-invert max-w-none prose-p:my-1 prose-pre:my-2 prose-headings:my-2 prose-ul:my-1 prose-ol:my-1 prose-li:my-0">
            <ReactMarkdown
              remarkPlugins={[remarkGfm]}
              components={{
                code({ className, children, ...props }) {
                  const match = /language-(\w+)/.exec(className || "")
                  const codeStr = String(children).replace(/\n$/, "")

                  if (match) {
                    return (
                      <div className="rounded border border-zinc-200 dark:border-zinc-700 overflow-hidden my-2 not-prose">
                        <CodeHighlighter
                          language={match[1]}
                          code={codeStr}
                          customStyle={CODE_STYLE}
                        />
                      </div>
                    )
                  }

                  return (
                    <code
                      className="px-1 py-0.5 rounded bg-zinc-200 dark:bg-zinc-700 text-xs font-mono"
                      {...props}
                    >
                      {children}
                    </code>
                  )
                },
              }}
            >
              {message.content}
            </ReactMarkdown>
          </div>
        ) : isStreaming ? (
          <Loader2 className="h-4 w-4 animate-spin text-zinc-400" />
        ) : null}
      </div>
      {isUser && (
        <div className="flex-shrink-0 w-7 h-7 rounded-full bg-zinc-200 dark:bg-zinc-700 flex items-center justify-center">
          <User className="h-3.5 w-3.5 text-zinc-600 dark:text-zinc-400" />
        </div>
      )}
    </div>
  )
})
