import { useEffect, useState } from 'react';
import { ArrowLeft, BarChart3, Wifi, Wrench, Layers } from 'lucide-react';
import { useNavigate } from 'react-router-dom';
import { api } from '@/api/client';
import type {
  AnalyticsSummary,
  SessionAnalyticsRow,
  ToolCostStat,
  ContextWindowStats,
  TokenBreakdownAgg,
} from '@/api/types';
import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card';
import { Button } from '@/components/ui/button';
import { cn } from '@/lib/utils';

// --- Independent section data fetcher ---
// Each section manages its own loading/error/data state so one failure doesn't
// block the rest of the page.

function useFetch<T>(fetcher: () => Promise<T>): {
  data: T | null;
  loading: boolean;
  error: string | null;
  retry: () => void;
} {
  const [data, setData] = useState<T | null>(null);
  const [loading, setLoading] = useState(true);
  const [error, setError] = useState<string | null>(null);

  function load() {
    setLoading(true);
    setError(null);
    fetcher()
      .then(setData)
      .catch((err: unknown) => {
        setError(err instanceof Error ? err.message : 'failed to load data');
      })
      .finally(() => setLoading(false));
  }

  useEffect(() => { load(); }, []); // eslint-disable-line react-hooks/exhaustive-deps

  return { data, loading, error, retry: load };
}

// --- Skeleton pulse placeholder ---
function SkeletonBar({ className }: { className?: string }) {
  return <div className={cn('animate-pulse rounded bg-muted/40', className)} />;
}

// --- Number formatting ---
function formatNumber(n: number | null | undefined): string {
  if (n == null) return '—';
  return n.toLocaleString();
}

function formatCost(n: number | null | undefined): string {
  if (n == null) return '—';
  if (n < 0.001) return `$${(n * 1000).toFixed(2)}m`;
  if (n < 0.01) return `$${n.toFixed(4)}`;
  return `$${n.toFixed(3)}`;
}

function formatPct(n: number | null | undefined): string {
  if (n == null) return '—';
  return `${(n * 100).toFixed(1)}%`;
}

function formatDate(iso: string | null | undefined): string {
  if (!iso) return '—';
  return new Date(iso).toLocaleDateString(undefined, {
    month: 'short',
    day: 'numeric',
    hour: '2-digit',
    minute: '2-digit',
  });
}

// --- Summary Cards ---
function SummaryCards({ summary }: { summary: AnalyticsSummary }) {
  const cards = [
    {
      label: 'Total Input Tokens',
      value: formatNumber(summary.total_input_tokens),
      icon: BarChart3,
      color: 'text-blue-500',
    },
    {
      label: 'Total Output Tokens',
      value: formatNumber(summary.total_output_tokens),
      icon: BarChart3,
      color: 'text-green-500',
    },
    {
      label: 'Total Cost',
      value: formatCost(summary.total_cost),
      icon: Wifi,
      color: 'text-amber-500',
    },
    {
      label: 'Sessions Tracked',
      value: formatNumber(summary.session_count),
      icon: Layers,
      color: 'text-purple-500',
    },
  ];

  return (
    <div className="grid grid-cols-2 md:grid-cols-4 gap-4">
      {cards.map((c) => (
        <Card key={c.label} size="sm">
          <CardContent className="flex items-start gap-3 pt-3">
            <c.icon className={cn('size-5 shrink-0 mt-0.5', c.color)} />
            <div className="min-w-0">
              <div className="text-lg font-semibold tabular-nums">{c.value}</div>
              <div className="text-xs text-muted-foreground mt-0.5">{c.label}</div>
            </div>
          </CardContent>
        </Card>
      ))}
    </div>
  );
}

function SummaryCardsSkeleton() {
  return (
    <div className="grid grid-cols-2 md:grid-cols-4 gap-4">
      {[0, 1, 2, 3].map((i) => (
        <Card key={i} size="sm">
          <CardContent className="pt-3">
            <SkeletonBar className="h-5 w-20 mb-2" />
            <SkeletonBar className="h-3 w-24" />
          </CardContent>
        </Card>
      ))}
    </div>
  );
}

// --- Section wrappers ---
function SectionCard({
  title,
  loading,
  error,
  onRetry,
  children,
}: {
  title: string;
  loading: boolean;
  error: string | null;
  onRetry: () => void;
  children: React.ReactNode;
}) {
  return (
    <Card>
      <CardHeader>
        <CardTitle>{title}</CardTitle>
      </CardHeader>
      <CardContent>
        {loading ? (
          <div className="space-y-2">
            <SkeletonBar className="h-4 w-full" />
            <SkeletonBar className="h-4 w-3/4" />
            <SkeletonBar className="h-4 w-1/2" />
          </div>
        ) : error ? (
          <div className="flex items-center gap-3 text-sm">
            <span className="text-destructive">{error}</span>
            <Button size="sm" variant="outline" onClick={onRetry}>
              Retry
            </Button>
          </div>
        ) : (
          children
        )}
      </CardContent>
    </Card>
  );
}

function EmptyState({ message }: { message: string }) {
  return <p className="text-sm text-muted-foreground py-2">{message}</p>;
}

// --- Per-Session Token Table ---
function SessionTable({ sessions }: { sessions: SessionAnalyticsRow[] }) {
  if (sessions.length === 0) {
    return <EmptyState message="No session token data available yet. Token data is collected as agent sessions run." />;
  }

  return (
    <div className="overflow-x-auto">
      <table className="w-full text-sm">
        <thead>
          <tr className="border-b text-left text-muted-foreground text-xs uppercase tracking-wide">
            <th className="py-2 pr-4 font-medium">Session</th>
            <th className="py-2 pr-4 font-medium tabular-nums text-right">Input</th>
            <th className="py-2 pr-4 font-medium tabular-nums text-right">Output</th>
            <th className="py-2 pr-4 font-medium tabular-nums text-right">Cost</th>
            <th className="py-2 font-medium tabular-nums text-right">Last Active</th>
          </tr>
        </thead>
        <tbody>
          {sessions.map((s) => (
            <tr key={s.session_id} className="border-b last:border-0 hover:bg-muted/30">
              <td className="py-2 pr-4 truncate max-w-[200px]" title={s.session_name}>
                {s.session_name || 'Untitled'}
              </td>
              <td className="py-2 pr-4 tabular-nums text-right">{formatNumber(s.total_input_tokens)}</td>
              <td className="py-2 pr-4 tabular-nums text-right">{formatNumber(s.total_output_tokens)}</td>
              <td className="py-2 pr-4 tabular-nums text-right">{formatCost(s.total_cost)}</td>
              <td className="py-2 tabular-nums text-right text-muted-foreground">{formatDate(s.last_active_at)}</td>
            </tr>
          ))}
        </tbody>
      </table>
    </div>
  );
}

// --- Per-Tool Cost Table ---
function ToolTable({ stats }: { stats: ToolCostStat[] }) {
  if (stats.length === 0) {
    return <EmptyState message="No tool cost data available yet. Stats accumulate after tool calls are made." />;
  }

  return (
    <div className="overflow-x-auto">
      <table className="w-full text-sm">
        <thead>
          <tr className="border-b text-left text-muted-foreground text-xs uppercase tracking-wide">
            <th className="py-2 pr-4 font-medium">Tool</th>
            <th className="py-2 pr-4 font-medium tabular-nums text-right">Calls</th>
            <th className="py-2 pr-4 font-medium tabular-nums text-right">Avg Prompt</th>
            <th className="py-2 pr-4 font-medium tabular-nums text-right">Avg Completion</th>
            <th className="py-2 font-medium tabular-nums text-right">Avg Total</th>
          </tr>
        </thead>
        <tbody>
          {stats.map((t) => (
            <tr key={t.tool_name} className="border-b last:border-0 hover:bg-muted/30">
              <td className="py-2 pr-4 flex items-center gap-2">
                <Wrench className="size-3.5 shrink-0 text-muted-foreground" />
                <span className="truncate max-w-[200px]" title={t.tool_name}>{t.tool_name}</span>
              </td>
              <td className="py-2 pr-4 tabular-nums text-right">{t.n_calls}</td>
              <td className="py-2 pr-4 tabular-nums text-right">{formatNumber(t.mean_prompt_tokens)}</td>
              <td className="py-2 pr-4 tabular-nums text-right">{formatNumber(t.mean_completion_tokens)}</td>
              <td className="py-2 tabular-nums text-right">{formatNumber(t.mean_prompt_tokens + t.mean_completion_tokens)}</td>
            </tr>
          ))}
        </tbody>
      </table>
    </div>
  );
}

// --- Context Window Utilization ---
function ContextSection({ stats }: { stats: ContextWindowStats }) {
  if (stats.message_count === 0) {
    return <EmptyState message="No context window data available yet. Data is captured during inference." />;
  }

  return (
    <div className="grid grid-cols-1 sm:grid-cols-3 gap-4">
      <div>
        <div className="text-xs text-muted-foreground">Avg Context Used</div>
        <div className="text-lg font-semibold tabular-nums mt-1">{formatNumber(Math.round(stats.avg_ctx_used ?? 0))}</div>
      </div>
      <div>
        <div className="text-xs text-muted-foreground">Avg Context Limit</div>
        <div className="text-lg font-semibold tabular-nums mt-1">{formatNumber(Math.round(stats.avg_ctx_max ?? 0))}</div>
      </div>
      <div>
        <div className="text-xs text-muted-foreground">Avg Utilization</div>
        <div className="text-lg font-semibold tabular-nums mt-1">{formatPct(stats.avg_utilization_pct)}</div>
      </div>
      <div className="sm:col-span-3">
        <div className="text-xs text-muted-foreground mb-1">Based on {formatNumber(stats.message_count)} completed assistant messages</div>
        <div className="h-2 rounded-full bg-muted overflow-hidden">
          <div
            className="h-full rounded-full bg-primary transition-all"
            style={{
              width: stats.avg_utilization_pct != null
                ? `${Math.min(stats.avg_utilization_pct * 100, 100)}%`
                : '0%',
            }}
          />
        </div>
      </div>
    </div>
  );
}

// --- Token Category Breakdown (CSS stacked bar) ---
const CATEGORY_COLORS: Record<string, string> = {
  system: 'bg-blue-500',
  user: 'bg-green-500',
  assistant: 'bg-amber-500',
  tools: 'bg-purple-500',
  reasoning: 'bg-rose-500',
};

const CATEGORY_LABELS: Record<string, string> = {
  system: 'System',
  user: 'User',
  assistant: 'Assistant',
  tools: 'Tools',
  reasoning: 'Reasoning',
};

function TokenBreakdownSection({ categories }: { categories: TokenBreakdownAgg[] }) {
  if (categories.length === 0) {
    return <EmptyState message="No token breakdown data available. Breakdown is captured for arena contestants and certain task types." />;
  }

  const total = categories.reduce((sum, c) => sum + c.total_tokens, 0);
  if (total === 0) return <EmptyState message="Token breakdown totals are zero." />;

  // Sort in a consistent order
  const order = ['system', 'user', 'assistant', 'tools', 'reasoning'];
  const sorted = [...categories].sort(
    (a, b) => order.indexOf(a.category) - order.indexOf(b.category),
  );

  return (
    <div className="space-y-3">
      <div className="h-4 rounded-full bg-muted overflow-hidden flex">
        {sorted.map((c) => {
          const pct = (c.total_tokens / total) * 100;
          if (pct < 1) return null;
          return (
            <div
              key={c.category}
              className={cn('h-full first:rounded-l-full last:rounded-r-full', CATEGORY_COLORS[c.category] ?? 'bg-gray-400')}
              style={{ width: `${pct}%` }}
              title={`${CATEGORY_LABELS[c.category] ?? c.category}: ${formatNumber(c.total_tokens)} (${pct.toFixed(1)}%)`}
            />
          );
        })}
      </div>
      <div className="flex flex-wrap gap-x-4 gap-y-1 text-xs">
        {sorted.map((c) => {
          const pct = (c.total_tokens / total) * 100;
          return (
            <div key={c.category} className="flex items-center gap-1.5">
              <span className={cn('size-2.5 rounded-sm', CATEGORY_COLORS[c.category] ?? 'bg-gray-400')} />
              <span className="text-muted-foreground">{CATEGORY_LABELS[c.category] ?? c.category}</span>
              <span className="font-medium tabular-nums">{pct.toFixed(1)}%</span>
              <span className="text-muted-foreground tabular-nums">({formatNumber(c.total_tokens)})</span>
            </div>
          );
        })}
      </div>
    </div>
  );
}

// --- Main Page ---
export function Analytics() {
  const navigate = useNavigate();

  const summary = useFetch(() => api.analytics.summary());
  const sessions = useFetch(() => api.analytics.sessions().then((r) => r.sessions));
  const tools = useFetch(() => api.tools.costStats().then((r) => r.stats));
  const context = useFetch(() => api.analytics.context());
  const breakdown = useFetch(() => api.analytics.tokenBreakdown().then((r) => r.categories));

  function handleBack() {
    if (window.history.length > 1) {
      navigate(-1);
    } else {
      navigate('/');
    }
  }

  return (
    <div className="flex-1 overflow-y-auto">
      <div className="max-w-[1000px] mx-auto w-full px-6 py-6 space-y-8">
        <header className="space-y-2">
          <button
            type="button"
            onClick={handleBack}
            className="inline-flex items-center gap-1 text-sm text-muted-foreground hover:text-foreground -ml-1 px-1 py-0.5 rounded"
            aria-label="Back"
          >
            <ArrowLeft className="size-4" />
            <span>Back</span>
          </button>
          <div>
            <h1 className="text-xl font-semibold">Token Analytics</h1>
            <p className="text-sm text-muted-foreground mt-1">
              Aggregate token usage, cost, and context window data across all sessions.
            </p>
          </div>
        </header>

        {/* Summary Cards */}
        {summary.loading ? (
          <SummaryCardsSkeleton />
        ) : summary.error ? (
          <div className="flex items-center gap-3 text-sm">
            <span className="text-destructive">{summary.error}</span>
            <Button size="sm" variant="outline" onClick={summary.retry}>
              Retry
            </Button>
          </div>
        ) : summary.data ? (
          <SummaryCards summary={summary.data} />
        ) : null}

        {/* Per-Session Token Breakdown */}
        <SectionCard
          title="Per-Session Token Usage"
          loading={sessions.loading}
          error={sessions.error}
          onRetry={sessions.retry}
        >
          {sessions.data && <SessionTable sessions={sessions.data} />}
        </SectionCard>

        {/* Per-Tool Cost Breakdown */}
        <SectionCard
          title="Per-Tool Token Cost"
          loading={tools.loading}
          error={tools.error}
          onRetry={tools.retry}
        >
          {tools.data && <ToolTable stats={tools.data} />}
        </SectionCard>

        {/* Context Window Utilization */}
        <SectionCard
          title="Context Window Utilization"
          loading={context.loading}
          error={context.error}
          onRetry={context.retry}
        >
          {context.data && <ContextSection stats={context.data} />}
        </SectionCard>

        {/* Token Category Breakdown */}
        <SectionCard
          title="Token Breakdown by Category"
          loading={breakdown.loading}
          error={breakdown.error}
          onRetry={breakdown.retry}
        >
          {breakdown.data && <TokenBreakdownSection categories={breakdown.data} />}
        </SectionCard>
      </div>
    </div>
  );
}
