import {
  Copy,
  Check,
  Sparkles,
  ClipboardList,
  Pencil,
  RefreshCw,
  GitFork,
  FolderGit2,
} from "lucide-react"
import { memo, useState, useMemo } from "react"
import { useNavigate } from "react-router-dom"
import { useCurrentUserQuery } from "@/hooks/useAuth"
import {
  useChatStreamRunner,
  useEditMessageMutation,
  useForkMessageMutation,
  useMessageVersionsQuery,
  useResendMessageMutation,
  useSwitchMessageVersionMutation,
} from "@/hooks/useAiChatQuery"
import { IrRenderer } from "@/lib/ir/renderer"
import {
  parseContentBlocks,
  extractAnswerText,
  extractFullText,
} from "@/lib/ir/parser"
import type {
  IrContentBlock,
  IrToolCallNode,
  IrToolResultNode,
} from "@/lib/ir/types"
import { Avatar, AvatarFallback, AvatarImage } from "@/components/ui/avatar"
import {
  Reasoning,
  ReasoningTrigger,
  ReasoningContent,
} from "@/components/ai-elements/reasoning"
import type { MessageResponse } from "@/hooks/useAiChatQuery"
import { getModelIcon } from "@/lib/icons/modelIcons"
import { ToolCallBlock } from "@/components/chat/ToolCallBlock"
import { Badge } from "@/components/ui/badge"
import { useChatPage } from "./ChatPageContext"
import { parseSlashContextMetadata } from "./chatSlashContext"

interface ChatMessageBubbleProps {
  message: MessageResponse
  conversationId: string
  onRegenerate?: (newMessageId: string) => void
  setIsStreaming: (value: boolean) => void
}

const AVATAR_COLORS = [
  "#6366f1",
  "#8b5cf6",
  "#d946ef",
  "#ec4899",
  "#f43f5e",
  "#ef4444",
  "#f97316",
  "#eab308",
  "#22c55e",
  "#14b8a6",
  "#06b6d4",
  "#3b82f6",
  "#2563eb",
  "#7c3aed",
  "#c026d3",
]

function hashColor(str: string): string {
  let hash = 0
  for (let i = 0; i < str.length; i++) {
    hash = str.charCodeAt(i) + ((hash << 5) - hash)
  }
  return AVATAR_COLORS[Math.abs(hash) % AVATAR_COLORS.length]
}

const PROSE_CLASS =
  "prose prose-sm dark:prose-invert max-w-none [&_p]:leading-[1.55] [&_p]:my-1 [&_ul]:my-1 [&_ol]:my-1 [&_li]:my-0 [&_h1]:mt-2 [&_h2]:mt-2 [&_h3]:mt-2 [&_pre]:my-1.5 [&_blockquote]:my-1"

export const ChatMessageBubble = memo(function ChatMessageBubble({
  message,
  conversationId,
  onRegenerate,
  setIsStreaming,
}: ChatMessageBubbleProps) {
  const isUser = message.role === "user"
  const [copied, setCopied] = useState<"answer" | "full" | false>(false)
  const [isEditing, setIsEditing] = useState(false)
  const [editText, setEditText] = useState("")
  const [showVersions, setShowVersions] = useState(false)
  const { data: user } = useCurrentUserQuery()
  const editMutation = useEditMessageMutation()
  const resendMutation = useResendMessageMutation()
  const forkMutation = useForkMessageMutation()
  const switchVersionMutation = useSwitchMessageVersionMutation()
  const runStream = useChatStreamRunner(setIsStreaming)
  const navigate = useNavigate()
  const { scope, projectName } = useChatPage()

  // Parse content into IrContentBlock[] (handles both old and future formats)
  const blocks: IrContentBlock[] = useMemo(
    () => (isUser ? [] : parseContentBlocks(message.content)),
    [isUser, message.content]
  )

  // User message plain text
  const userText =
    typeof message.content === "string"
      ? message.content
      : message.content && typeof message.content === "object"
        ? String((message.content as Record<string, unknown>).content ?? "")
        : ""
  const plainText = isUser ? userText : extractAnswerText(blocks)
  const hasThinking = blocks.some((b) => b.role === "thinking")
  const selectedContexts = useMemo(
    () => parseSlashContextMetadata(message.metadata),
    [message.metadata]
  )

  // Fetch versions when showing version switcher
  const versionsQuery = useMessageVersionsQuery(conversationId, message.id)

  const handleCopyAnswer = () => {
    const text = isUser ? plainText : extractAnswerText(blocks)
    if (text) {
      navigator.clipboard.writeText(text)
      setCopied("answer")
      setTimeout(() => setCopied(false), 2000)
    }
  }

  const handleCopyFull = () => {
    const text = isUser ? plainText : extractFullText(blocks)
    if (text) {
      navigator.clipboard.writeText(text)
      setCopied("full")
      setTimeout(() => setCopied(false), 2000)
    }
  }

  const handleStartEdit = () => {
    setEditText(plainText)
    setIsEditing(true)
  }

  const handleCancelEdit = () => {
    setIsEditing(false)
    setEditText("")
  }

  const handleSaveEdit = async () => {
    if (!editText.trim()) return
    try {
      const newMsg = await editMutation.mutateAsync({
        conversationId,
        messageId: message.id,
        content: editText.trim(),
      })
      setIsEditing(false)
      onRegenerate?.(newMsg.id)
      await runStream(conversationId, newMsg.id)
    } catch (err) {
      console.error("Failed to edit message:", err)
    }
  }

  const handleRegenerate = async () => {
    try {
      const newMsg = await resendMutation.mutateAsync({
        conversationId,
        messageId: message.id,
      })
      onRegenerate?.(newMsg.id)
      await runStream(conversationId, newMsg.id)
    } catch (err) {
      console.error("Failed to regenerate:", err)
    }
  }

  const handleSwitchVersion = async (versionNumber: number) => {
    try {
      await switchVersionMutation.mutateAsync({
        conversationId,
        messageId: message.id,
        versionNumber,
      })
      setShowVersions(false)
    } catch (err) {
      console.error("Failed to switch version:", err)
    }
  }

  const handleFork = async () => {
    try {
      const fork = await forkMutation.mutateAsync({
        conversationId,
        messageId: message.id,
      })
      if (scope === "project" && projectName) {
        navigate(`/${projectName}/chat/${fork.id}`)
      } else {
        navigate(`/me/chat/${fork.id}`)
      }
    } catch (err) {
      console.error("Failed to fork:", err)
    }
  }

  return (
    <div className="group mx-auto flex w-full max-w-3xl gap-4 rounded-2xl border border-transparent px-4 py-4">
      {/* Avatar */}
      <div className="shrink-0 pt-0.5">
        {isUser ? (
          <Avatar className="h-7 w-7 rounded-full">
            <AvatarImage src={user?.avatar_url ?? undefined} />
            <AvatarFallback
              className="rounded-full text-[10px] font-semibold"
              style={{
                backgroundColor: hashColor(user?.username || "user"),
                color: "var(--text-inverse)",
              }}
            >
              {(user?.display_name || user?.username || "U")[0]?.toUpperCase()}
            </AvatarFallback>
          </Avatar>
        ) : (
          <ModelAvatar modelName={message.model} size={28} />
        )}
      </div>

      {/* Content */}
      <div className="min-w-0 flex-1">
        {/* Sender name */}
        <div className="mb-2 flex flex-wrap items-center gap-2">
          <span
            className="text-xs font-semibold"
            style={{ color: "var(--text-primary)" }}
          >
            {isUser
              ? user?.display_name || user?.username || "You"
              : message.model || "Assistant"}
          </span>
          <span className="text-[11px]" style={{ color: "var(--text-muted)" }}>
            {new Date(message.created_at).toLocaleTimeString([], {
              hour: "2-digit",
              minute: "2-digit",
            })}
          </span>
          {!isUser && message.version_number > 1 && (
            <Badge
              variant="outline"
              className="rounded-full px-2 py-0.5 text-[10px] tracking-[0.16em] uppercase"
            >
              v{message.version_number}
            </Badge>
          )}
        </div>

        {/* Interleaved blocks — thinking (collapsible) + answer (IrRenderer) */}
        <div className="text-sm" style={{ color: "var(--text-primary)" }}>
          {isUser && selectedContexts.length > 0 && (
            <div className="mb-2 flex flex-wrap gap-2">
              {selectedContexts.map((context) => (
                <Badge
                  key={`${context.kind}:${context.id}`}
                  variant="outline"
                  className="h-auto gap-1 rounded-full px-2.5 py-1"
                >
                  {context.kind === "repo" ? (
                    <FolderGit2 className="h-3 w-3" />
                  ) : (
                    <Sparkles className="h-3 w-3" />
                  )}
                  <span>{context.label}</span>
                  <span className="text-[10px] uppercase opacity-70">
                    {context.kind}
                  </span>
                </Badge>
              ))}
            </div>
          )}

          {isUser && isEditing ? (
            <div className="flex flex-col gap-2 rounded-2xl border border-[var(--border-subtle)] bg-[var(--surface-ground)] p-3">
              <textarea
                value={editText}
                onChange={(e) => setEditText(e.target.value)}
                className="w-full resize-none rounded-xl p-3 text-[14px] outline-none"
                style={{
                  backgroundColor: "var(--surface-elevated)",
                  border: "1px solid var(--accent)",
                  color: "var(--text-primary)",
                  lineHeight: "1.6",
                  minHeight: "60px",
                }}
                autoFocus
                rows={3}
                onKeyDown={(e) => {
                  if (e.key === "Enter" && !e.shiftKey) {
                    e.preventDefault()
                    handleSaveEdit()
                  }
                  if (e.key === "Escape") {
                    handleCancelEdit()
                  }
                }}
              />
              <div className="flex items-center justify-end gap-2">
                <button
                  onClick={handleCancelEdit}
                  className="rounded-full px-3 py-1.5 text-[12px] transition-colors"
                  style={{
                    color: "var(--text-muted)",
                    border: "1px solid var(--border-default)",
                  }}
                >
                  Cancel
                </button>
                <button
                  onClick={handleSaveEdit}
                  disabled={!editText.trim() || editMutation.isPending}
                  className="rounded-full px-3 py-1.5 text-[12px] font-medium transition-colors"
                  style={{
                    backgroundColor: editText.trim()
                      ? "var(--accent)"
                      : "var(--muted)",
                    color: editText.trim()
                      ? "var(--text-inverse)"
                      : "var(--text-muted)",
                  }}
                >
                  {editMutation.isPending ? "Saving..." : "Save & Submit"}
                </button>
              </div>
            </div>
          ) : isUser ? (
            <p className="whitespace-pre-wrap">{plainText}</p>
          ) : (
            blocks.map((b, i) => {
              if (b.role === "thinking") {
                // Thinking content rendered by Reasoning/Streamdown (not IrRenderer)
                const thinkingText = b.nodes
                  .filter((n) => n.type === "text")
                  .map((n) => (n as { content: string }).content)
                  .join("")
                return (
                  <Reasoning key={i} defaultOpen={false}>
                    <ReasoningTrigger />
                    <ReasoningContent>{thinkingText}</ReasoningContent>
                  </Reasoning>
                )
              }
              if (b.role === "tool_call") {
                // Tool call visualization
                const toolCallNode = b.nodes.find(
                  (n) => n.type === "tool_call"
                ) as IrToolCallNode | undefined
                if (toolCallNode && toolCallNode.tool === "call_sub_agent") {
                  return (
                    <ToolCallBlock
                      key={i}
                      toolName={toolCallNode.tool}
                      args={toolCallNode.args}
                      status="ok"
                      conversationId={conversationId}
                    />
                  )
                }
                return null
              }
              if (b.role === "tool_result") {
                const toolResultNode = b.nodes.find(
                  (n) => n.type === "tool_result"
                ) as IrToolResultNode | undefined
                if (
                  toolResultNode &&
                  toolResultNode.tool === "call_sub_agent"
                ) {
                  return (
                    <ToolCallBlock
                      key={i}
                      toolName={toolResultNode.tool}
                      args={{
                        role: toolResultNode.role,
                        task: toolResultNode.task,
                      }}
                      status={toolResultNode.status}
                      result={toolResultNode.content}
                      childrenId={toolResultNode.children_id}
                      subAgentOutput={toolResultNode.content}
                      conversationId={conversationId}
                    />
                  )
                }
                return null
              }
              return (
                <div key={i} className={i > 0 ? "mt-3" : ""}>
                  <IrRenderer nodes={b.nodes} className={PROSE_CLASS} />
                </div>
              )
            })
          )}
        </div>

        {/* Version indicator */}
        {message.version_number > 1 && (
          <button
            onClick={() => setShowVersions(!showVersions)}
            className="mt-2 inline-flex items-center gap-2 rounded-full border border-[var(--border-subtle)] bg-[var(--surface-secondary)] px-2.5 py-1 text-[11px] text-muted-foreground"
          >
            <span className="font-medium text-[var(--text-primary)]">
              v{message.version_number}
            </span>
            <span>{showVersions ? "Hide versions" : "View all versions"}</span>
          </button>
        )}

        {/* Version switcher */}
        {showVersions &&
          versionsQuery.data &&
          versionsQuery.data.length > 1 && (
            <div className="mt-2 rounded-2xl border border-[var(--border-subtle)] bg-[var(--surface-elevated)] p-2">
              {versionsQuery.data.map((v) => (
                <button
                  key={v.id}
                  onClick={() => handleSwitchVersion(v.version_number)}
                  className="flex w-full items-center gap-2 rounded-xl px-3 py-2 text-left"
                  style={{
                    color:
                      v.version_number === message.version_number
                        ? "var(--accent)"
                        : "var(--text-muted)",
                    fontSize: "12px",
                  }}
                >
                  <Badge
                    variant={
                      v.version_number === message.version_number
                        ? "secondary"
                        : "outline"
                    }
                    className="rounded-full px-2 py-0.5 text-[10px]"
                  >
                    v{v.version_number}
                  </Badge>
                  <span className="flex-1 truncate">
                    {typeof v.content === "string"
                      ? v.content.slice(0, 50)
                      : String(v.content).slice(0, 50)}
                  </span>
                  <span className="text-[11px] opacity-60">
                    {new Date(v.created_at).toLocaleTimeString([], {
                      hour: "2-digit",
                      minute: "2-digit",
                    })}
                  </span>
                </button>
              ))}
            </div>
          )}

        {/* User message actions — edit & regenerate */}
        {isUser && !isEditing && (
          <div className="mt-2 flex flex-wrap items-center gap-1 opacity-0 transition-opacity group-hover:opacity-100">
            <button
              onClick={handleStartEdit}
              className="inline-flex items-center gap-1 rounded-full px-2.5 py-1 text-[11px]"
              style={{ color: "var(--text-muted)" }}
            >
              <Pencil className="size-3" />
              Edit
            </button>
            {onRegenerate && (
              <button
                onClick={handleRegenerate}
                className="inline-flex items-center gap-1 rounded-full px-2.5 py-1 text-[11px]"
                style={{ color: "var(--text-muted)" }}
              >
                <RefreshCw className="size-3" />
                Regenerate
              </button>
            )}
          </div>
        )}

        {/* AI message actions — copy & fork */}
        {!isUser && (
          <div className="mt-2 flex flex-wrap items-center gap-1 opacity-0 transition-opacity group-hover:opacity-100">
            <button
              onClick={handleCopyAnswer}
              className="inline-flex items-center gap-1 rounded-full px-2.5 py-1 text-[11px]"
              style={{ color: "var(--text-muted)" }}
            >
              {copied === "answer" ? (
                <Check className="size-3" />
              ) : (
                <Copy className="size-3" />
              )}
              {copied === "answer" ? "Copied" : "Copy"}
            </button>
            {hasThinking && (
              <button
                onClick={handleCopyFull}
                className="inline-flex items-center gap-1 rounded-full px-2.5 py-1 text-[11px]"
                style={{ color: "var(--text-muted)" }}
              >
                {copied === "full" ? (
                  <Check className="size-3" />
                ) : (
                  <ClipboardList className="size-3" />
                )}
                {copied === "full" ? "Copied" : "Copy with reasoning"}
              </button>
            )}
            <button
              onClick={handleFork}
              className="inline-flex items-center gap-1 rounded-full px-2.5 py-1 text-[11px]"
              style={{ color: "var(--text-muted)" }}
            >
              <GitFork className="size-3" />
              Fork
            </button>
          </div>
        )}

        {/* Token usage */}
        {!isUser &&
          (message.input_tokens != null || message.output_tokens != null) && (
            <div className="mt-2 flex flex-wrap gap-1.5">
              {message.input_tokens != null && (
                <Badge
                  variant="outline"
                  className="rounded-full px-2 py-0.5 text-[10px]"
                >
                  {message.input_tokens} in
                </Badge>
              )}
              {message.output_tokens != null && (
                <Badge
                  variant="outline"
                  className="rounded-full px-2 py-0.5 text-[10px]"
                >
                  {message.output_tokens} out
                </Badge>
              )}
              {(message.input_tokens ?? 0) + (message.output_tokens ?? 0) >
                0 && (
                <Badge
                  variant="secondary"
                  className="rounded-full px-2 py-0.5 text-[10px]"
                >
                  {(message.input_tokens ?? 0) + (message.output_tokens ?? 0)}{" "}
                  total
                </Badge>
              )}
            </div>
          )}
      </div>
    </div>
  )
})

// ─── Model Avatar ─────────────────────────────────────────────────────

function ModelAvatar({
  modelName,
  size = 28,
}: {
  modelName?: string | null
  size?: number
}) {
  if (!modelName) {
    return (
      <div
        className="flex h-7 w-7 shrink-0 items-center justify-center rounded-full"
        style={{
          width: size,
          height: size,
          backgroundColor: "var(--accent-muted)",
        }}
      >
        <Sparkles
          className="h-3.5 w-3.5"
          style={{
            color: "var(--accent)",
            width: size * 0.5,
            height: size * 0.5,
          }}
        />
      </div>
    )
  }
  const Icon = getModelIcon(modelName)
  if (Icon) {
    return (
      <div className="shrink-0" style={{ width: size, height: size }}>
        <Icon.Avatar size={size} />
      </div>
    )
  }
  return (
    <div
      className="flex shrink-0 items-center justify-center rounded-full font-bold"
      style={{
        width: size,
        height: size,
        backgroundColor: hashColor(modelName),
        fontSize: Math.max(10, size * 0.35),
        color: "var(--text-inverse)",
      }}
    >
      {modelName[0]?.toUpperCase() || "?"}
    </div>
  )
}