'use client';

import { useState, useEffect, useCallback } from 'react';
import {
  Sparkles,
  FileText,
  Zap,
  Activity,
  RefreshCw,
  ChevronDown,
  ChevronUp,
  Tag,
  Clock,
} from 'lucide-react';
import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card';
import { Badge } from '@/components/ui/badge';
import { Button } from '@/components/ui/button';
import {
  Table,
  TableBody,
  TableCell,
  TableHead,
  TableHeader,
  TableRow,
} from '@/components/ui/table';
import { Separator } from '@/components/ui/separator';
import { apiFetch } from '@/lib/api';

// ── Types ────────────────────────────────────────────────────────

interface ExtractionEntity {
  extraction_class: string;
  extraction_text: string;
  attributes?: Record<string, string>;
}

interface ExtractResponse {
  extractions: ExtractionEntity[];
  metadata: {
    modelId: string;
    durationMs: number;
    tokenCount?: number;
    charCount: number;
  };
  requestId?: string;
}

interface ExtractionTask {
  id: string;
  name: string;
  description?: string;
  prompt: string;
  classes: string[];
  builtIn: boolean;
  productId: string;
}

interface SidecarHealth {
  status: string;
  sidecar?: { status: string; version: string };
}

// ── Helpers ──────────────────────────────────────────────────────

const CLASS_COLORS: Record<string, string> = {
  action_item: 'bg-blue-500/20 text-blue-400 border-blue-500/30',
  decision: 'bg-purple-500/20 text-purple-400 border-purple-500/30',
  question: 'bg-yellow-500/20 text-yellow-400 border-yellow-500/30',
  deadline: 'bg-red-500/20 text-red-400 border-red-500/30',
  person: 'bg-green-500/20 text-green-400 border-green-500/30',
  topic: 'bg-cyan-500/20 text-cyan-400 border-cyan-500/30',
  emotion: 'bg-pink-500/20 text-pink-400 border-pink-500/30',
  brain_signal: 'bg-indigo-500/20 text-indigo-400 border-indigo-500/30',
  entity: 'bg-teal-500/20 text-teal-400 border-teal-500/30',
  action: 'bg-orange-500/20 text-orange-400 border-orange-500/30',
};

function getClassColor(cls: string): string {
  return CLASS_COLORS[cls] || 'bg-zinc-500/20 text-zinc-400 border-zinc-500/30';
}

// ── Component ────────────────────────────────────────────────────

export default function ExtractionPage() {
  const [inputText, setInputText] = useState('');
  const [selectedTask, setSelectedTask] = useState('transcript-extraction');
  const [tasks, setTasks] = useState<ExtractionTask[]>([]);
  const [result, setResult] = useState<ExtractResponse | null>(null);
  const [health, setHealth] = useState<SidecarHealth | null>(null);
  const [loading, setLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [showTaskDetails, setShowTaskDetails] = useState(false);

  // Load tasks and health on mount
  useEffect(() => {
    loadTasks();
    checkHealth();
  }, []);

  async function loadTasks() {
    try {
      const res = await apiFetch<ExtractionTask[]>('/extraction/tasks');
      if (res.data) setTasks(res.data);
    } catch {
      // Tasks endpoint may not be available yet
    }
  }

  async function checkHealth() {
    try {
      const res = await apiFetch<SidecarHealth>('/extraction/extract/sidecar-health');
      setHealth(res.data);
    } catch {
      setHealth({ status: 'error' });
    }
  }

  const handleExtract = useCallback(async () => {
    if (!inputText.trim()) return;
    setLoading(true);
    setError(null);
    setResult(null);

    try {
      const res = await apiFetch<ExtractResponse>('/extraction/extract', {
        method: 'POST',
        body: JSON.stringify({
          text: inputText,
          taskId: selectedTask,
        }),
      });
      if (res.data) setResult(res.data);
      else setError(res.error || 'No response from extraction service');
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Extraction failed');
    } finally {
      setLoading(false);
    }
  }, [inputText, selectedTask]);

  const currentTask = tasks.find(t => t.id === selectedTask);

  // Group extractions by class
  const groupedExtractions = result?.extractions.reduce(
    (acc, e) => {
      const cls = e.extraction_class;
      if (!acc[cls]) acc[cls] = [];
      acc[cls].push(e);
      return acc;
    },
    {} as Record<string, ExtractionEntity[]>
  );

  return (
    <div className="space-y-6">
      {/* Header */}
      <div className="flex items-center justify-between">
        <div>
          <h1 className="text-2xl font-bold tracking-tight">Extraction</h1>
          <p className="text-muted-foreground text-sm">
            Extract structured entities from text using LangExtract
          </p>
        </div>
        <div className="flex items-center gap-2">
          <Badge
            variant="outline"
            className={
              health?.status === 'ok'
                ? 'border-green-500/50 text-green-400'
                : 'border-red-500/50 text-red-400'
            }
          >
            <Activity className="mr-1 h-3 w-3" />
            {health?.status === 'ok' ? 'Sidecar Online' : 'Sidecar Offline'}
          </Badge>
          <Button variant="ghost" size="icon" onClick={checkHealth}>
            <RefreshCw className="h-4 w-4" />
          </Button>
        </div>
      </div>

      <Separator />

      {/* Input Section */}
      <div className="grid gap-4 md:grid-cols-3">
        <div className="md:col-span-2 space-y-3">
          <label className="text-sm font-medium">Input Text</label>
          <textarea
            className="w-full min-h-[200px] rounded-lg border border-border bg-background p-3 text-sm resize-y focus:outline-none focus:ring-2 focus:ring-ring"
            placeholder="Paste a transcript, meeting notes, or any text to extract structured entities from..."
            value={inputText}
            onChange={e => setInputText(e.target.value)}
          />
          <div className="flex items-center gap-2">
            <Button onClick={handleExtract} disabled={loading || !inputText.trim()}>
              {loading ? (
                <RefreshCw className="mr-2 h-4 w-4 animate-spin" />
              ) : (
                <Sparkles className="mr-2 h-4 w-4" />
              )}
              {loading ? 'Extracting...' : 'Extract'}
            </Button>
            <span className="text-xs text-muted-foreground">
              {inputText.length.toLocaleString()} chars
            </span>
          </div>
        </div>

        {/* Task Selector */}
        <div className="space-y-3">
          <label className="text-sm font-medium">Task</label>
          <select
            className="w-full rounded-lg border border-border bg-background p-2 text-sm"
            value={selectedTask}
            onChange={e => setSelectedTask(e.target.value)}
          >
            {tasks.length > 0 ? (
              tasks.map(t => (
                <option key={t.id} value={t.id}>
                  {t.name} {t.builtIn ? '(built-in)' : ''}
                </option>
              ))
            ) : (
              <>
                <option value="transcript-extraction">Transcript Extraction</option>
                <option value="triage">MindLyst Triage</option>
                <option value="memory-insight">Memory Insight</option>
                <option value="reflection-enrichment">Reflection Enrichment</option>
                <option value="bug-report-extraction">Bug Report Extraction</option>
              </>
            )}
          </select>

          {currentTask && (
            <div className="space-y-2">
              <button
                className="flex items-center text-xs text-muted-foreground hover:text-foreground"
                onClick={() => setShowTaskDetails(!showTaskDetails)}
              >
                {showTaskDetails ? (
                  <ChevronUp className="mr-1 h-3 w-3" />
                ) : (
                  <ChevronDown className="mr-1 h-3 w-3" />
                )}
                Task Details
              </button>
              {showTaskDetails && (
                <Card>
                  <CardContent className="pt-4 space-y-2 text-xs">
                    <p className="text-muted-foreground">{currentTask.description}</p>
                    <div className="flex flex-wrap gap-1">
                      {currentTask.classes.map(cls => (
                        <Badge
                          key={cls}
                          variant="outline"
                          className={`text-[10px] ${getClassColor(cls)}`}
                        >
                          {cls}
                        </Badge>
                      ))}
                    </div>
                  </CardContent>
                </Card>
              )}
            </div>
          )}

          {/* Quick Stats */}
          <Card>
            <CardContent className="pt-4 space-y-2">
              <div className="flex items-center justify-between text-xs">
                <span className="text-muted-foreground">Model</span>
                <span>gemini-2.5-flash</span>
              </div>
              <div className="flex items-center justify-between text-xs">
                <span className="text-muted-foreground">Tasks</span>
                <span>{tasks.length || 5} available</span>
              </div>
            </CardContent>
          </Card>
        </div>
      </div>

      {/* Error */}
      {error && (
        <Card className="border-red-500/30 bg-red-500/5">
          <CardContent className="pt-4">
            <p className="text-sm text-red-400">{error}</p>
          </CardContent>
        </Card>
      )}

      {/* Results */}
      {result && (
        <div className="space-y-4">
          <Separator />

          {/* Result Metadata */}
          <div className="grid gap-4 md:grid-cols-4">
            <Card>
              <CardHeader className="flex flex-row items-center justify-between space-y-0 pb-2">
                <CardTitle className="text-sm font-medium">Entities</CardTitle>
                <Tag className="h-4 w-4 text-muted-foreground" />
              </CardHeader>
              <CardContent>
                <div className="text-2xl font-bold">{result.extractions.length}</div>
              </CardContent>
            </Card>
            <Card>
              <CardHeader className="flex flex-row items-center justify-between space-y-0 pb-2">
                <CardTitle className="text-sm font-medium">Duration</CardTitle>
                <Clock className="h-4 w-4 text-muted-foreground" />
              </CardHeader>
              <CardContent>
                <div className="text-2xl font-bold">{result.metadata.durationMs}ms</div>
              </CardContent>
            </Card>
            <Card>
              <CardHeader className="flex flex-row items-center justify-between space-y-0 pb-2">
                <CardTitle className="text-sm font-medium">Model</CardTitle>
                <Zap className="h-4 w-4 text-muted-foreground" />
              </CardHeader>
              <CardContent>
                <div className="text-lg font-bold truncate">{result.metadata.modelId}</div>
              </CardContent>
            </Card>
            <Card>
              <CardHeader className="flex flex-row items-center justify-between space-y-0 pb-2">
                <CardTitle className="text-sm font-medium">Input</CardTitle>
                <FileText className="h-4 w-4 text-muted-foreground" />
              </CardHeader>
              <CardContent>
                <div className="text-2xl font-bold">
                  {result.metadata.charCount.toLocaleString()}
                </div>
                <p className="text-xs text-muted-foreground">characters</p>
              </CardContent>
            </Card>
          </div>

          {/* Grouped Extractions */}
          {groupedExtractions && Object.keys(groupedExtractions).length > 0 && (
            <div className="space-y-3">
              <h3 className="text-sm font-medium">Extracted Entities by Class</h3>
              <div className="grid gap-3 md:grid-cols-2">
                {Object.entries(groupedExtractions).map(([cls, entities]) => (
                  <Card key={cls}>
                    <CardHeader className="pb-2">
                      <div className="flex items-center justify-between">
                        <Badge variant="outline" className={getClassColor(cls)}>
                          {cls}
                        </Badge>
                        <span className="text-xs text-muted-foreground">
                          {entities.length} found
                        </span>
                      </div>
                    </CardHeader>
                    <CardContent className="space-y-2">
                      {entities.map((e, i) => (
                        <div
                          key={i}
                          className="rounded-md border border-border/50 bg-muted/30 p-2 text-sm"
                        >
                          <p className="text-foreground">&quot;{e.extraction_text}&quot;</p>
                          {e.attributes && Object.keys(e.attributes).length > 0 && (
                            <div className="mt-1 flex flex-wrap gap-1">
                              {Object.entries(e.attributes).map(([k, v]) => (
                                <Badge key={k} variant="secondary" className="text-[10px]">
                                  {k}: {v}
                                </Badge>
                              ))}
                            </div>
                          )}
                        </div>
                      ))}
                    </CardContent>
                  </Card>
                ))}
              </div>
            </div>
          )}

          {/* Raw Results Table */}
          <div className="space-y-3">
            <h3 className="text-sm font-medium">All Extractions</h3>
            <Card>
              <Table>
                <TableHeader>
                  <TableRow>
                    <TableHead className="w-[150px]">Class</TableHead>
                    <TableHead>Text</TableHead>
                    <TableHead className="w-[200px]">Attributes</TableHead>
                  </TableRow>
                </TableHeader>
                <TableBody>
                  {result.extractions.map((e, i) => (
                    <TableRow key={i}>
                      <TableCell>
                        <Badge
                          variant="outline"
                          className={`text-xs ${getClassColor(e.extraction_class)}`}
                        >
                          {e.extraction_class}
                        </Badge>
                      </TableCell>
                      <TableCell className="text-sm">{e.extraction_text}</TableCell>
                      <TableCell className="text-xs text-muted-foreground">
                        {e.attributes
                          ? Object.entries(e.attributes)
                              .map(([k, v]) => `${k}=${v}`)
                              .join(', ')
                          : '—'}
                      </TableCell>
                    </TableRow>
                  ))}
                </TableBody>
              </Table>
            </Card>
          </div>
        </div>
      )}
    </div>
  );
}
