"use client";

import { useMemo } from "react";
import { Activity, DollarSign } from "lucide-react";
import { cn, formatDuration } from "@/lib/utils";

interface DecisionPoint {
  id: string;
  type: string;
  chosenAction: string;
  alternatives: string[];
  reasoning: string | null;
  contextSnapshot: Record<string, unknown> | null;
  confidence: number | null;
  timestamp: string;
  parentSpanId?: string | null;
}

interface Span {
  id: string;
  name: string;
  type: string;
  status: "OK" | "ERROR" | "CANCELLED";
  startedAt: string;
  endedAt: string | null;
  durationMs: number | null;
  input: unknown;
  output: unknown;
  metadata: Record<string, unknown>;
  parentSpanId?: string | null;
}

interface Event {
  id: string;
  type: string;
  name: string;
  timestamp: string;
  metadata: Record<string, unknown>;
  spanId?: string | null;
}

interface Trace {
  id: string;
  name: string;
  status: "RUNNING" | "COMPLETED" | "ERROR";
  startedAt: string;
  endedAt: string | null;
  durationMs: number | null;
  tags: string[];
  metadata: Record<string, unknown>;
  costUsd: number | null;
  totalTokens: number | null;
  totalCost: number | null;
}

interface TraceAnalyticsProps {
  trace: Trace;
  spans: Span[];
  decisionPoints: DecisionPoint[];
  events: Event[];
}

const spanTypeColors: Record<string, { bg: string; text: string }> = {
  LLM_CALL: { bg: "bg-purple-500/10", text: "text-purple-400" },
  TOOL_CALL: { bg: "bg-blue-500/10", text: "text-blue-400" },
  DEFAULT: { bg: "bg-neutral-700/30", text: "text-neutral-400" },
};

const statusBarColors: Record<string, string> = {
  OK: "bg-emerald-500",
  ERROR: "bg-red-500",
  CANCELLED: "bg-amber-500",
};

export function TraceAnalytics({
  trace,
  spans,
  decisionPoints,
}: TraceAnalyticsProps) {
  return (
    <div className="space-y-6">
      {/* Section A: Execution Timeline */}
      <section className="p-5 bg-neutral-900 border border-neutral-800 rounded-xl">
        <h2 className="text-sm font-semibold text-neutral-300 uppercase tracking-wider mb-4">
          Execution Timeline
        </h2>
        <ExecutionTimeline trace={trace} spans={spans} />
      </section>

      {/* Section B: Cost Breakdown */}
      <section className="p-5 bg-neutral-900 border border-neutral-800 rounded-xl">
        <h2 className="text-sm font-semibold text-neutral-300 uppercase tracking-wider mb-4">
          Cost Breakdown
        </h2>
        <CostBreakdown trace={trace} spans={spans} decisionPoints={decisionPoints} />
      </section>

      {/* Section C: Token Usage Gauge */}
      <section className="p-5 bg-neutral-900 border border-neutral-800 rounded-xl">
        <h2 className="text-sm font-semibold text-neutral-300 uppercase tracking-wider mb-4">
          Token Usage
        </h2>
        <TokenUsageGauge trace={trace} />
      </section>
    </div>
  );
}

// Section A: Execution Timeline (Waterfall Chart)
function ExecutionTimeline({ trace, spans }: { trace: Trace; spans: Span[] }) {
  const timelineData = useMemo(() => {
    if (!trace.startedAt || spans.length === 0) {
      return null;
    }

    const traceStartTime = new Date(trace.startedAt).getTime();
    const lastSpanEnd = spans.reduce((max, s) => {
      const end = s.endedAt ? new Date(s.endedAt).getTime() : 0;
      return end > max ? end : max;
    }, 0);
    const traceEndTime = trace.endedAt
      ? new Date(trace.endedAt).getTime()
      : trace.status === "COMPLETED" || trace.status === "ERROR"
        ? lastSpanEnd || traceStartTime
        : Date.now();
    const totalDuration = traceEndTime - traceStartTime;

    // Build span hierarchy for nesting
    const spanMap = new Map<string, Span & { depth: number }>();
    const rootSpans: (Span & { depth: number })[] = [];

    // Sort spans by start time first
    const sortedSpans = [...spans].sort(
      (a, b) => new Date(a.startedAt).getTime() - new Date(b.startedAt).getTime()
    );

    // Calculate depth for each span
    sortedSpans.forEach((span) => {
      let depth = 0;
      let parentId = span.parentSpanId;
      while (parentId) {
        depth++;
        const parent = spans.find((s) => s.id === parentId);
        parentId = parent?.parentSpanId;
      }
      const spanWithDepth = { ...span, depth };
      spanMap.set(span.id, spanWithDepth);
      if (!span.parentSpanId) {
        rootSpans.push(spanWithDepth);
      }
    });

    // Build nested display order
    const displaySpans: (Span & { depth: number })[] = [];
    const addSpanAndChildren = (spanId: string) => {
      const span = spanMap.get(spanId);
      if (span) {
        displaySpans.push(span);
        // Find children
        const children = sortedSpans.filter((s) => s.parentSpanId === spanId);
        children.forEach((child) => addSpanAndChildren(child.id));
      }
    };
    rootSpans.forEach((root) => addSpanAndChildren(root.id));

    return {
      traceStartTime,
      totalDuration,
      displaySpans,
    };
  }, [trace, spans]);

  if (!timelineData) {
    return (
      <div className="text-center py-8 text-neutral-500">
        No spans to visualize
      </div>
    );
  }

  const { traceStartTime, totalDuration, displaySpans } = timelineData;

  // Generate time markers
  const timeMarkers = useMemo(() => {
    const markers: { label: string; percent: number }[] = [];
    const step = totalDuration <= 1000 ? 100 : totalDuration <= 5000 ? 500 : 1000;
    const count = Math.ceil(totalDuration / step);

    for (let i = 0; i <= count; i++) {
      const timeMs = i * step;
      if (timeMs <= totalDuration) {
        const percent = (timeMs / totalDuration) * 100;
        let label: string;
        if (timeMs < 1000) {
          label = `${timeMs}ms`;
        } else {
          label = `${(timeMs / 1000).toFixed(timeMs % 1000 === 0 ? 0 : 1)}s`;
        }
        markers.push({ label, percent });
      }
    }
    return markers;
  }, [totalDuration]);

  return (
    <div className="space-y-4">
      {/* Ruler */}
      <div className="relative h-6 border-b border-neutral-700 ml-[200px]">
        {timeMarkers.map((marker, idx) => (
          <div
            key={idx}
            className="absolute top-0 transform -translate-x-1/2"
            style={{ left: `${marker.percent}%` }}
          >
            <div className="w-px h-2 bg-neutral-600" />
            <span className="text-xs text-neutral-500 absolute top-3 -translate-x-1/2 whitespace-nowrap">
              {marker.label}
            </span>
          </div>
        ))}
      </div>

      {/* Span rows */}
      <div className="space-y-0.5">
        {displaySpans.map((span) => {
          const spanStartTime = new Date(span.startedAt).getTime();
          const offsetPercent =
            ((spanStartTime - traceStartTime) / totalDuration) * 100;
          const durationPercent =
            ((span.durationMs || 0) / totalDuration) * 100;
          const isRunning = !span.endedAt;
          const typeColors = spanTypeColors[span.type] || spanTypeColors.DEFAULT;
          const statusColor = statusBarColors[span.status] || statusBarColors.OK;

          return (
            <div
              key={span.id}
              className="flex items-center h-9"
              style={{ marginLeft: `${span.depth * 16}px` }}
            >
              {/* Span info (left column) */}
              <div className="w-[200px] flex items-center gap-2 pr-4 shrink-0">
                <span className="text-sm text-neutral-300 truncate">
                  {span.name.length > 20 ? `${span.name.slice(0, 20)}...` : span.name}
                </span>
                <span
                  className={cn(
                    "px-1.5 py-0.5 rounded text-[10px] font-medium uppercase tracking-wide shrink-0",
                    typeColors.bg,
                    typeColors.text
                  )}
                >
                  {span.type}
                </span>
              </div>

              {/* Timeline bar area */}
              <div className="flex-1 relative h-full">
                <div
                  className={cn(
                    "absolute top-2 h-5 rounded",
                    statusColor,
                    isRunning && "animate-pulse"
                  )}
                  style={{
                    left: `${Math.max(0, offsetPercent)}%`,
                    width: `${Math.max(1, durationPercent)}%`,
                    minWidth: isRunning ? "4px" : "1px",
                  }}
                />
              </div>
            </div>
          );
        })}
      </div>

      {/* Duration info */}
      <div className="text-xs text-neutral-500 mt-2 pt-2 border-t border-neutral-800">
        Total Duration: {formatDuration(totalDuration)}
      </div>
    </div>
  );
}

// Section B: Cost Breakdown
function CostBreakdown({
  trace,
  spans,
  decisionPoints,
}: {
  trace: Trace;
  spans: Span[];
  decisionPoints: DecisionPoint[];
}) {
  const { spanCosts, decisionCosts, totalCost, hasCostData } = useMemo(() => {
    const spanCostsData = spans
      .map((span) => ({
        id: span.id,
        name: span.name,
        type: span.type,
        cost: (span.metadata?.costUsd as number | null | undefined) ?? null,
      }))
      .filter((s) => s.cost !== null && s.cost > 0)
      .sort((a, b) => (b.cost || 0) - (a.cost || 0));

    // Decision costs - use null since costUsd may not be in the API response yet
    const decisionCostsData = decisionPoints
      .map((dp) => ({
        id: dp.id,
        name: dp.chosenAction,
        type: dp.type,
        cost: null as number | null,
      }))
      .filter((d) => d.cost !== null && d.cost > 0);

    const totalSpanCost = spanCostsData.reduce((sum, s) => sum + (s.cost || 0), 0);
    const totalDecisionCost = decisionCostsData.reduce(
      (sum, d) => sum + (d.cost || 0),
      0
    );
    const totalCostValue = trace.costUsd ?? trace.totalCost ?? totalSpanCost + totalDecisionCost;
    const hasData =
      totalCostValue > 0 || spanCostsData.length > 0 || decisionCostsData.length > 0;

    const maxSpanCost =
      spanCostsData.length > 0 ? Math.max(...spanCostsData.map((s) => s.cost || 0)) : 0;
    const maxDecisionCost =
      decisionCostsData.length > 0
        ? Math.max(...decisionCostsData.map((d) => d.cost || 0))
        : 0;

    return {
      spanCosts: spanCostsData.map((s) => ({
        ...s,
        percent: maxSpanCost > 0 ? ((s.cost || 0) / maxSpanCost) * 100 : 0,
      })),
      decisionCosts: decisionCostsData.map((d) => ({
        ...d,
        percent: maxDecisionCost > 0 ? ((d.cost || 0) / maxDecisionCost) * 100 : 0,
      })),
      totalCost: totalCostValue,
      hasCostData: hasData,
    };
  }, [trace, spans, decisionPoints]);

  if (!hasCostData) {
    return (
      <div className="text-center py-8 text-neutral-500">
        <DollarSign className="w-8 h-8 mx-auto mb-2 text-neutral-600" />
        No cost data available
      </div>
    );
  }

  return (
    <div className="space-y-4">
      {/* Total cost summary */}
      <div className="flex items-center gap-2 pb-4 border-b border-neutral-800">
        <DollarSign className="w-5 h-5 text-emerald-400" />
        <span className="text-lg font-semibold text-neutral-100">
          Total: ${totalCost.toFixed(4)}
        </span>
      </div>

      {/* Cost breakdown columns */}
      <div className="grid grid-cols-1 lg:grid-cols-2 gap-6">
        {/* Span Costs */}
        <div>
          <h3 className="text-xs font-medium text-neutral-500 uppercase tracking-wide mb-3">
            Span Costs
          </h3>
          <div className="space-y-2">
            {spanCosts.length > 0 ? (
              spanCosts.map((span) => {
                const typeColors =
                  spanTypeColors[span.type] || spanTypeColors.DEFAULT;
                return (
                  <div key={span.id} className="space-y-1">
                    <div className="flex items-center justify-between text-sm">
                      <div className="flex items-center gap-2">
                        <span
                          className={cn(
                            "px-1.5 py-0.5 rounded text-[10px] font-medium uppercase",
                            typeColors.bg,
                            typeColors.text
                          )}
                        >
                          {span.type}
                        </span>
                        <span className="text-neutral-300 truncate max-w-[150px]">
                          {span.name.length > 20
                            ? `${span.name.slice(0, 20)}...`
                            : span.name}
                        </span>
                      </div>
                      <span className="text-emerald-400 font-medium">
                        ${(span.cost || 0).toFixed(4)}
                      </span>
                    </div>
                    <div className="h-2 bg-neutral-800 rounded-full overflow-hidden">
                      <div
                        className="h-full bg-emerald-500 rounded-full transition-all duration-300"
                        style={{ width: `${span.percent}%` }}
                      />
                    </div>
                  </div>
                );
              })
            ) : (
              <p className="text-sm text-neutral-500">No span cost data</p>
            )}
          </div>
        </div>

        {/* Decision Costs */}
        <div>
          <h3 className="text-xs font-medium text-neutral-500 uppercase tracking-wide mb-3">
            Decision Costs
          </h3>
          <div className="space-y-2">
            {decisionCosts.length > 0 ? (
              decisionCosts.map((decision) => {
                const typeColors =
                  spanTypeColors[decision.type] || spanTypeColors.DEFAULT;
                return (
                  <div key={decision.id} className="space-y-1">
                    <div className="flex items-center justify-between text-sm">
                      <div className="flex items-center gap-2">
                        <span
                          className={cn(
                            "px-1.5 py-0.5 rounded text-[10px] font-medium uppercase",
                            typeColors.bg,
                            typeColors.text
                          )}
                        >
                          {decision.type}
                        </span>
                        <span className="text-neutral-300 truncate max-w-[150px]">
                          {decision.name.length > 20
                            ? `${decision.name.slice(0, 20)}...`
                            : decision.name}
                        </span>
                      </div>
                      <span className="text-emerald-400 font-medium">
                        {decision.cost !== null
                          ? `$${decision.cost.toFixed(4)}`
                          : "N/A"}
                      </span>
                    </div>
                    <div className="h-2 bg-neutral-800 rounded-full overflow-hidden">
                      <div
                        className="h-full bg-emerald-500 rounded-full transition-all duration-300"
                        style={{ width: `${decision.percent}%` }}
                      />
                    </div>
                  </div>
                );
              })
            ) : (
              <p className="text-sm text-neutral-500">No decision cost data</p>
            )}
          </div>
        </div>
      </div>
    </div>
  );
}

// Section C: Token Usage Gauge
function TokenUsageGauge({ trace }: { trace: Trace }) {
  const tokenData = useMemo(() => {
    const totalTokens =
      trace.totalTokens ??
      (trace.metadata?.totalTokens as number | null | undefined) ??
      (trace.metadata?.tokenCount as number | null | undefined) ??
      null;

    const modelContextWindows: Record<string, number> = {
      "gpt-5.2": 128000,
      "gpt-5.1": 128000,
      "gpt-5": 128000,
      "gpt-5-mini": 128000,
      "gpt-5-nano": 128000,
      "gpt-4.1": 1047576,
      "gpt-4.1-mini": 1047576,
      "gpt-4.1-nano": 1047576,
      "o3": 200000,
      "o3-mini": 200000,
      "o4-mini": 200000,
      "gpt-4": 8192,
      "gpt-4-32k": 32768,
      "gpt-4-turbo": 128000,
      "gpt-4o": 128000,
      "gpt-4o-mini": 128000,
      "gpt-3.5-turbo": 16385,
      "claude-opus-4-6": 200000,
      "claude-4.5-opus": 200000,
      "claude-4.5-sonnet": 200000,
      "claude-4.5-haiku": 200000,
      "claude-3-opus": 200000,
      "claude-3-sonnet": 200000,
      "claude-3-haiku": 200000,
      "claude-3.5-sonnet": 200000,
      "claude-4-opus": 200000,
      "claude-4-sonnet": 200000,
    };

    const model = (trace.metadata?.model as string | undefined) ?? "";
    const modelLower = model.toLowerCase();
    let maxTokens = 128000;
    for (const [prefix, ctx] of Object.entries(modelContextWindows)) {
      if (modelLower.startsWith(prefix)) {
        maxTokens = ctx;
        break;
      }
    }

    return {
      totalTokens,
      maxTokens,
      percent: totalTokens !== null ? (totalTokens / maxTokens) * 100 : 0,
      hasData: totalTokens !== null && totalTokens > 0,
    };
  }, [trace]);

  const { totalTokens, maxTokens, percent, hasData } = tokenData;

  // Determine color based on usage percentage
  const getUsageColor = (pct: number): string => {
    if (pct <= 50) return "#10b981"; // emerald-500
    if (pct <= 80) return "#f59e0b"; // amber-500
    return "#ef4444"; // red-500
  };

  const usageColor = getUsageColor(percent);

  // Create conic gradient for the ring
  const ringStyle = {
    background: hasData
      ? `conic-gradient(${usageColor} ${percent * 3.6}deg, #262626 ${percent * 3.6}deg)`
      : "conic-gradient(#404040 0deg, #404040 360deg)",
  };

  return (
    <div className="flex flex-col items-center justify-center py-6">
      {/* Ring container */}
      <div className="relative">
        {/* Outer ring with conic gradient */}
        <div
          className="w-40 h-40 rounded-full p-3"
          style={ringStyle}
        >
          {/* Inner circle (background) */}
          <div className="w-full h-full rounded-full bg-neutral-900 flex items-center justify-center">
            {/* Center content */}
            <div className="text-center">
              {hasData ? (
                <>
                  <div className="text-3xl font-bold text-neutral-100">
                    {totalTokens?.toLocaleString()}
                  </div>
                  <div className="text-xs text-neutral-500 mt-1">tokens</div>
                </>
              ) : (
                <Activity className="w-8 h-8 text-neutral-600 mx-auto" />
              )}
            </div>
          </div>
        </div>

        {/* Percentage indicator */}
        {hasData && (
          <div
            className="absolute -top-1 left-1/2 transform -translate-x-1/2 w-8 h-8 rounded-full flex items-center justify-center text-xs font-bold"
            style={{ backgroundColor: usageColor }}
          >
            <span className="text-neutral-950">{Math.round(percent)}%</span>
          </div>
        )}
      </div>

      {/* Token info below */}
      <div className="mt-6 text-center space-y-1">
        {hasData ? (
          <>
            <p className="text-lg font-medium text-neutral-200">
              {totalTokens?.toLocaleString()} / {maxTokens.toLocaleString()} tokens
            </p>
            <p className="text-sm text-neutral-500">
              {percent.toFixed(1)}% of context window used
            </p>
            {percent > 80 && (
              <p className="text-xs text-red-400 mt-2">
                High token usage - consider context optimization
              </p>
            )}
          </>
        ) : (
          <>
            <p className="text-lg font-medium text-neutral-400">No token data</p>
            <p className="text-sm text-neutral-600">
              Token usage information not available for this trace
            </p>
          </>
        )}
      </div>
    </div>
  );
}