import { useCallback, useEffect, useMemo, useState } from 'react';
import { ChevronDown, ChevronRight, AlertCircle } from 'lucide-react';
import { api } from '@/api/client';
import type { ToolTrace } from '@/api/types';

interface Props {
  chatId: string;
}

// Max latency used as the 100% reference for the bar visualization
const MAX_LATENCY_REF = 30_000; // 30s

function latencyBarWidth(latencyMs: number | null): number {
  if (latencyMs == null) return 0;
  return Math.min(latencyMs / MAX_LATENCY_REF, 1);
}

function TraceRow({ trace }: { trace: ToolTrace }) {
  const [expanded, setExpanded] = useState(false);
  const isError = trace.outcome !== null && trace.outcome !== 'success';
  const barWidth = latencyBarWidth(trace.latency_ms);
  const latencyLabel =
    trace.latency_ms != null
      ? trace.latency_ms >= 1000
        ? `${(trace.latency_ms / 1000).toFixed(1)}s`
        : `${trace.latency_ms}ms`
      : null;

  return (
    <div className="border-b border-border/40 last:border-0">
      <button
        type="button"
        onClick={() => setExpanded((v) => !v)}
        className="flex items-center gap-2 w-full text-left px-2 py-1.5 hover:bg-muted/40 text-[11px]"
      >
        <span className="shrink-0 text-muted-foreground">
          {expanded ? <ChevronDown size={10} /> : <ChevronRight size={10} />}
        </span>
        <span className="font-medium truncate min-w-0">
          {trace.tool_name}
        </span>
        {isError && (
          <span className="shrink-0 text-destructive" title={trace.error ?? 'error'}>
            <AlertCircle size={10} />
          </span>
        )}
        <span className="shrink-0 text-muted-foreground font-mono tabular-nums min-w-[3rem] text-right">
          {latencyLabel ?? '—'}
        </span>
        <span className="flex-1 h-1.5 bg-muted rounded-full overflow-hidden min-w-[24px] max-w-[60px]">
          <span
            className="block h-full rounded-full bg-primary/30 transition-all"
            style={{ width: `${barWidth * 100}%` }}
          />
        </span>
        {trace.tokens_used != null && trace.tokens_used > 0 && (
          <span className="shrink-0 text-muted-foreground font-mono tabular-nums">
            {trace.tokens_used}t
          </span>
        )}
        {trace.cache_tokens != null && trace.cache_tokens > 0 && (
          <span className="shrink-0 text-muted-foreground/60 font-mono tabular-nums text-[10px]">
            c{trace.cache_tokens}
          </span>
        )}
        {trace.reasoning_tokens != null && trace.reasoning_tokens > 0 && (
          <span className="shrink-0 text-muted-foreground/60 font-mono tabular-nums text-[10px]">
            r{trace.reasoning_tokens}
          </span>
        )}
      </button>
      {expanded && (
        <div className="px-3 pb-2 space-y-1.5 text-[11px] border-t border-border/40 pt-1.5">
          <div>
            <span className="text-muted-foreground font-medium">Input</span>
            <pre className="mt-0.5 font-mono text-[10px] leading-relaxed text-muted-foreground bg-muted/30 rounded p-1.5 overflow-x-auto max-h-32 overflow-y-auto whitespace-pre-wrap break-all">
              {JSON.stringify(trace.tool_input, null, 1)}
            </pre>
          </div>
          {trace.tool_output != null && (
            <div>
              <span className="text-muted-foreground font-medium">Output</span>
              <pre className="mt-0.5 font-mono text-[10px] leading-relaxed text-muted-foreground bg-muted/30 rounded p-1.5 overflow-x-auto max-h-32 overflow-y-auto whitespace-pre-wrap break-all">
                {trace.tool_output.length > 2000
                  ? `${trace.tool_output.slice(0, 2000)}…`
                  : trace.tool_output}
              </pre>
            </div>
          )}
          {trace.error != null && (
            <div className="text-destructive text-[10px] font-mono leading-relaxed bg-destructive/10 rounded p-1.5">
              {trace.error}
            </div>
          )}
        </div>
      )}
    </div>
  );
}

function TraceGroup({ toolName, traces }: { toolName: string; traces: ToolTrace[] }) {
  const [collapsed, setCollapsed] = useState(false);
  const totalLatency = traces.reduce((sum, t) => sum + (t.latency_ms ?? 0), 0);
  const totalTokens = traces.reduce((sum, t) => sum + (t.tokens_used ?? 0), 0);
  const errorCount = traces.filter(
    (t) => t.outcome !== null && t.outcome !== 'success',
  ).length;

  return (
    <div>
      <button
        type="button"
        onClick={() => setCollapsed((v) => !v)}
        className="flex items-center gap-1.5 w-full text-left px-2 py-1 text-[11px] font-medium text-muted-foreground hover:bg-muted/30 sticky top-0 bg-background"
      >
        {collapsed ? <ChevronRight size={10} /> : <ChevronDown size={10} />}
        <span>{toolName}</span>
        <span className="text-muted-foreground/60 font-mono tabular-nums">
          ×{traces.length}
        </span>
        {totalTokens > 0 && (
          <span className="text-muted-foreground/60 font-mono tabular-nums text-[10px]">
            {totalTokens}t
          </span>
        )}
        {totalLatency > 0 && (
          <span className="text-muted-foreground/60 font-mono tabular-nums text-[10px]">
            {totalLatency >= 1000
              ? `${(totalLatency / 1000).toFixed(1)}s`
              : `${totalLatency}ms`}
          </span>
        )}
        {errorCount > 0 && (
          <span className="ml-auto text-destructive text-[10px] font-medium">
            {errorCount} error{errorCount > 1 ? 's' : ''}
          </span>
        )}
      </button>
      {!collapsed && traces.map((trace) => (
        <TraceRow key={trace.id} trace={trace} />
      ))}
    </div>
  );
}

export function TraceViewer({ chatId }: Props) {
  const [open, setOpen] = useState(false);
  const [traces, setTraces] = useState<ToolTrace[]>([]);
  const [loading, setLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);

  const fetchTraces = useCallback(async () => {
    setLoading(true);
    setError(null);
    try {
      const res = await api.chats.getTraces(chatId);
      setTraces(res.data);
    } catch (err) {
      setError(err instanceof Error ? err.message : 'failed to load traces');
    } finally {
      setLoading(false);
    }
  }, [chatId]);

  useEffect(() => {
    if (open) {
      void fetchTraces();
    }
  }, [open, fetchTraces]);

  const groups = useMemo(() => {
    const map = new Map<string, ToolTrace[]>();
    for (const t of traces) {
      const existing = map.get(t.tool_name);
      if (existing) {
        existing.push(t);
      } else {
        map.set(t.tool_name, [t]);
      }
    }
    return map;
  }, [traces]);

  const totalCount = traces.length;
  const errorCount = traces.filter(
    (t) => t.outcome !== null && t.outcome !== 'success',
  ).length;

  return (
    <div className="border-t">
      <button
        type="button"
        onClick={() => setOpen((v) => !v)}
        className="flex items-center gap-1.5 w-full px-3 py-1.5 text-[11px] font-medium text-muted-foreground hover:bg-muted/20"
      >
        {open ? <ChevronDown size={12} /> : <ChevronRight size={12} />}
        <span>Tool traces</span>
        {totalCount > 0 && (
          <span className="font-mono tabular-nums text-muted-foreground/60">
            {totalCount}
          </span>
        )}
        {errorCount > 0 && (
          <span className="text-destructive ml-auto text-[10px] font-medium">
            {errorCount} error{errorCount > 1 ? 's' : ''}
          </span>
        )}
        {loading && (
          <span className="ml-auto inline-block w-1.5 h-3 align-baseline bg-muted-foreground/60 animate-pulse" />
        )}
      </button>
      {open && (
        <div className="max-h-80 overflow-y-auto border-t border-border/40">
          {loading && traces.length === 0 && (
            <div className="px-3 py-4 text-[11px] text-muted-foreground text-center">
              Loading traces…
            </div>
          )}
          {error && (
            <div className="px-3 py-2 text-[11px] text-destructive">
              {error}
              <button
                type="button"
                onClick={() => void fetchTraces()}
                className="ml-2 underline hover:no-underline"
              >
                retry
              </button>
            </div>
          )}
          {!loading && !error && traces.length === 0 && (
            <div className="px-3 py-4 text-[11px] text-muted-foreground text-center">
              No tool traces yet.
            </div>
          )}
          {traces.length > 0 && (
            <div className="divide-y divide-border/40">
              {Array.from(groups.entries()).map(([toolName, groupTraces]) => (
                <TraceGroup
                  key={toolName}
                  toolName={toolName}
                  traces={groupTraces}
                />
              ))}
            </div>
          )}
        </div>
      )}
    </div>
  );
}