"use client";

import * as React from "react";
import ReactMarkdown from "react-markdown";
import remarkGfm from "remark-gfm";
import { CodeBlock } from "@/components/chat/code-block";
import { ThinkingBlock } from "@/components/chat/thinking-block";
import { ToolCallBlock, type ToolCallView } from "@/components/chat/tool-call-block";
import { Avatar, AvatarFallback } from "@/components/ui/avatar";
import { Button } from "@/components/ui/button";
import { Check, Copy, RotateCw, Clock, Gauge, Hash, Loader2 } from "lucide-react";
import { useSession } from "next-auth/react";
import { cn } from "@/lib/utils";
import type { MessageRecord } from "@/components/chat/chat-view";

interface Props {
  message: MessageRecord;
  showThinking: boolean;
  streaming?: boolean;
  streamingStartedAt?: number;
  streamingLiveTokens?: number;
  streamingToolCalls?: ToolCallView[];
  onRegenerate?: () => void;
}

export function Message({
  message,
  showThinking,
  streaming,
  streamingStartedAt,
  streamingLiveTokens,
  streamingToolCalls,
  onRegenerate,
}: Props) {
  const { data: session } = useSession();
  const isUser = message.role === "user";
  const isTool = message.role === "tool";
  const [copied, setCopied] = React.useState(false);

  // Tool result messages: handled inside their preceding assistant message's tool blocks.
  if (isTool) return null;

  async function onCopy() {
    try {
      await navigator.clipboard.writeText(message.content);
      setCopied(true);
      setTimeout(() => setCopied(false), 1500);
    } catch {
      /* ignore */
    }
  }

  const userInitial = (session?.user?.username ?? "U")[0]?.toUpperCase() ?? "U";

  const hasFinalStats =
    !isUser &&
    (message.tokens !== null ||
      message.evalDurationMs !== null ||
      message.totalDurationMs !== null);

  return (
    <div className={cn("group mb-6 flex gap-4", isUser ? "" : "")}>
      <div className="shrink-0">
        <Avatar className="h-7 w-7">
          <AvatarFallback
            className={cn(
              "text-xs font-medium",
              isUser ? "bg-muted text-foreground" : "bg-primary text-primary-foreground",
            )}
          >
            {isUser ? userInitial : "AI"}
          </AvatarFallback>
        </Avatar>
      </div>
      <div className="min-w-0 flex-1">
        <div className="mb-1 flex items-center gap-2">
          <span className="text-xs font-medium text-muted-foreground">
            {isUser ? session?.user?.username ?? "You" : "Assistant"}
          </span>
        </div>
        {showThinking && (message.thinking || (streaming && message.thinking !== null)) && (
          <ThinkingBlock
            thinking={message.thinking ?? ""}
            streaming={streaming && (!message.content || message.content.length === 0)}
          />
        )}

        {streaming && streamingToolCalls && streamingToolCalls.length > 0 && (
          <div className="mb-2">
            {streamingToolCalls.map((tc) => (
              <ToolCallBlock key={tc.id} call={tc} />
            ))}
          </div>
        )}

        {!streaming && !isUser && <PersistedToolCalls message={message} />}

        {isUser ? (
          <div className="whitespace-pre-wrap text-[0.975rem] leading-7">{message.content}</div>
        ) : message.content ? (
          <div className="prose-chat">
            <ReactMarkdown
              remarkPlugins={[remarkGfm]}
              components={{
                code({ className, children, ...props }) {
                  const text = String(children ?? "").replace(/\n$/, "");
                  const match = /language-(\w+)/.exec(className || "");
                  const isInline = !match && !text.includes("\n");
                  if (isInline) {
                    return (
                      <code className={className} {...props}>
                        {children}
                      </code>
                    );
                  }
                  return <CodeBlock code={text} language={match?.[1]} />;
                },
              }}
            >
              {message.content || (streaming ? "​" : "")}
            </ReactMarkdown>
            {streaming && message.content && <span className="caret-blink" />}
          </div>
        ) : null}

        {!isUser && streaming && !hasFinalStats && (
          <LiveStats
            startedAt={streamingStartedAt ?? Date.now()}
            liveTokens={streamingLiveTokens ?? 0}
          />
        )}
        {hasFinalStats && <FinalStats message={message} />}

        {!streaming && message.content && (
          <div className="mt-2 flex items-center gap-1 opacity-0 transition-opacity group-hover:opacity-100">
            <Button variant="ghost" size="sm" className="h-7 gap-1 px-2 text-xs" onClick={onCopy}>
              {copied ? <Check className="!h-3 !w-3" /> : <Copy className="!h-3 !w-3" />}
              {copied ? "Copied" : "Copy"}
            </Button>
            {onRegenerate && (
              <Button
                variant="ghost"
                size="sm"
                className="h-7 gap-1 px-2 text-xs"
                onClick={onRegenerate}
              >
                <RotateCw className="!h-3 !w-3" />
                Regenerate
              </Button>
            )}
          </div>
        )}
      </div>
    </div>
  );
}

/**
 * Renders persisted tool calls: parses `message.toolCalls` (the calls emitted
 * by this assistant message) and pairs them with their result messages stored
 * in the chat. The result messages are looked up via React context.
 */
function PersistedToolCalls({ message }: { message: MessageRecord }) {
  const results = useToolResultsFor(message);
  if (!message.toolCalls) return null;
  let calls: Array<{ id?: string; function?: { name?: string; arguments?: unknown } }> = [];
  try {
    calls = JSON.parse(message.toolCalls);
  } catch {
    return null;
  }
  if (!Array.isArray(calls) || calls.length === 0) return null;

  return (
    <div className="mb-2">
      {calls.map((c, i) => {
        const id = c.id ?? `${message.id}-${i}`;
        const name = c.function?.name ?? "tool";
        const args = (c.function?.arguments ?? {}) as Record<string, unknown> | string;
        const match = results.find((r) => r.toolCallId === id || r.toolName === name);
        return (
          <ToolCallBlock
            key={id}
            call={{
              id,
              name,
              arguments: args,
              status: match ? "done" : "done",
              result: match?.content,
              meta: match?.toolMeta ? parseMeta(match.toolMeta) : null,
              durationMs: match?.totalDurationMs ?? undefined,
            }}
          />
        );
      })}
    </div>
  );
}

function parseMeta(s: string): Record<string, unknown> | null {
  try {
    const v = JSON.parse(s);
    return v && typeof v === "object" ? v : null;
  } catch {
    return null;
  }
}

// Lightweight lookup: we read sibling messages from the parent chat via context.
import { ChatMessagesContext } from "@/components/chat/chat-messages-context";

function useToolResultsFor(message: MessageRecord): MessageRecord[] {
  const all = React.useContext(ChatMessagesContext);
  if (!all) return [];
  const idx = all.findIndex((m) => m.id === message.id);
  if (idx === -1) return [];
  // Tool results follow the assistant message, until the next non-tool message.
  const out: MessageRecord[] = [];
  for (let i = idx + 1; i < all.length; i++) {
    const m = all[i];
    if (!m) break;
    if (m.role !== "tool") break;
    out.push(m);
  }
  return out;
}

function LiveStats({ startedAt, liveTokens }: { startedAt: number; liveTokens: number }) {
  const [now, setNow] = React.useState(() => Date.now());
  React.useEffect(() => {
    const id = setInterval(() => setNow(Date.now()), 100);
    return () => clearInterval(id);
  }, []);
  const elapsedSec = Math.max(0, (now - startedAt) / 1000);
  const tps = elapsedSec > 0.2 && liveTokens > 0 ? liveTokens / elapsedSec : null;

  return (
    <StatsPill streaming>
      <Loader2 className="h-3 w-3 animate-spin" />
      <StatItem icon={<Hash className="h-3 w-3" />} label={`${liveTokens} tok`} />
      <StatDivider />
      <StatItem icon={<Clock className="h-3 w-3" />} label={`${elapsedSec.toFixed(1)}s`} />
      {tps !== null && (
        <>
          <StatDivider />
          <StatItem icon={<Gauge className="h-3 w-3" />} label={`${tps.toFixed(1)} tok/s`} />
        </>
      )}
    </StatsPill>
  );
}

function FinalStats({ message }: { message: MessageRecord }) {
  const tokens = message.tokens;
  const promptTokens = message.promptTokens;
  const evalMs = message.evalDurationMs;
  const totalMs = message.totalDurationMs;

  const tokPerSec = tokens && evalMs && evalMs > 0 ? tokens / (evalMs / 1000) : null;

  const seconds = totalMs ?? evalMs;
  const secLabel = seconds
    ? seconds >= 1000
      ? `${(seconds / 1000).toFixed(2)}s`
      : `${seconds}ms`
    : null;

  return (
    <StatsPill>
      {tokens !== null && (
        <StatItem
          icon={<Hash className="h-3 w-3" />}
          label={
            <>
              {tokens.toLocaleString()} tok
              {promptTokens !== null && (
                <span className="opacity-60"> · {promptTokens.toLocaleString()} in</span>
              )}
            </>
          }
        />
      )}
      {secLabel && (
        <>
          {tokens !== null && <StatDivider />}
          <StatItem icon={<Clock className="h-3 w-3" />} label={secLabel} />
        </>
      )}
      {tokPerSec !== null && (
        <>
          <StatDivider />
          <StatItem icon={<Gauge className="h-3 w-3" />} label={`${tokPerSec.toFixed(1)} tok/s`} />
        </>
      )}
    </StatsPill>
  );
}

function StatsPill({ children, streaming }: { children: React.ReactNode; streaming?: boolean }) {
  return (
    <div
      className={cn(
        "mt-3 inline-flex flex-wrap items-center gap-2 rounded-full border border-border bg-muted/40 px-3 py-1 text-[11px] font-medium text-muted-foreground",
        streaming && "border-primary/30 bg-primary/5 text-primary",
      )}
    >
      {children}
    </div>
  );
}

function StatItem({ icon, label }: { icon: React.ReactNode; label: React.ReactNode }) {
  return (
    <span className="inline-flex items-center gap-1 tabular-nums">
      {icon}
      {label}
    </span>
  );
}

function StatDivider() {
  return <span className="opacity-40">·</span>;
}