src/app/components/ModelExplorer.tsx

import { useState, useEffect } from 'react';
import { Card } from './ui/card';
import { Input } from './ui/input';
import { Button } from './ui/button';
import { Badge } from './ui/badge';
import { BarChart, Bar, XAxis, YAxis, Tooltip, ResponsiveContainer } from 'recharts';
import { Database, Search, ChevronLeft, ChevronRight, Settings, Cpu, LineChart, Sparkles } from 'lucide-react';
import { API_BASE_URL } from '../config';

interface ModelMetrics {
  modelName: string;
  datasetSize: number;
  numClasses: number;
  accuracy: number;
  macroF1: number;
  weightedF1: number;
  weightedPrecision: number;
  weightedRecall: number;
  vocabularySize: number;
  wordFeatureCount: number;
  charFeatureCount: number;
  ngramRange: [number, number];
  charNgramRange: [number, number];
  maxIterations: number;
  splitMethod: string;
  groupOverlap: number;
  calibrated: boolean;
}

export function ModelExplorer() {
  // Search & Filter State
  const [searchQuery, setSearchQuery] = useState('');
  const [selectedCategory, setSelectedCategory] = useState('');
  const [currentPage, setCurrentPage] = useState(0);
  const [totalRecords, setTotalRecords] = useState(0);
  const [records, setRecords] = useState<any[]>([]);
  const [categoryCounts, setCategoryCounts] = useState<Record<string, number>>({});
  const [metrics, setMetrics] = useState<ModelMetrics | null>(null);
  
  // loading state
  const [isLoading, setIsLoading] = useState(true);
  const [isTableLoading, setIsTableLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);

  const PAGE_LIMIT = 10;

  // Load category counts initially
  useEffect(() => {
    async function loadStats() {
      setIsLoading(true);
      try {
        const [datasetResponse, metricsResponse] = await Promise.all([
          fetch(`${API_BASE_URL}/api/dataset?limit=1`),
          fetch(`${API_BASE_URL}/api/metrics`),
        ]);
        if (!datasetResponse.ok || !metricsResponse.ok) {
          throw new Error('Failed to load model statistics');
        }
        const [data, metricsData] = await Promise.all([
          datasetResponse.json(),
          metricsResponse.json(),
        ]);
        if (data.status === 'success') {
          setCategoryCounts(data.categoryCounts);
        }
        setMetrics(metricsData);
      } catch (err: any) {
        console.error(err);
        setError(err.message || 'Failed to connect to dataset API.');
      } finally {
        setIsLoading(false);
      }
    }
    loadStats();
  }, []);

  // Query records when page/search/category changes
  useEffect(() => {
    async function fetchRecords() {
      setIsTableLoading(true);
      const offset = currentPage * PAGE_LIMIT;
      const url = `${API_BASE_URL}/api/dataset?q=${encodeURIComponent(searchQuery)}&category=${encodeURIComponent(selectedCategory)}&limit=${PAGE_LIMIT}&offset=${offset}`;
      
      try {
        const res = await fetch(url);
        if (!res.ok) throw new Error('Failed to fetch records');
        const data = await res.json();
        if (data.status === 'success') {
          setRecords(data.records);
          setTotalRecords(data.total);
        }
      } catch (err: any) {
        console.error(err);
      } finally {
        setIsTableLoading(false);
      }
    }
    
    // Debounce search input
    const delayDebounceFn = setTimeout(() => {
      fetchRecords();
    }, searchQuery ? 300 : 0);

    return () => clearTimeout(delayDebounceFn);
  }, [searchQuery, selectedCategory, currentPage]);

  const handleSearchChange = (e: React.ChangeEvent<HTMLInputElement>) => {
    setSearchQuery(e.target.value);
    setCurrentPage(0); // Reset page on new query
  };

  const handleCategoryChange = (category: string) => {
    setSelectedCategory(category);
    setCurrentPage(0);
  };

  // Convert category stats to Recharts format
  const chartData = Object.entries(categoryCounts)
    .map(([name, value]) => ({ name, count: value }))
    .sort((a, b) => b.count - a.count);

  const getSeverityBadgeClass = (category: string) => {
    switch (category.toLowerCase()) {
      case 'critical':
      case 'forced action':
      case 'sneaking':
      case 'obstruction':
        return 'bg-rose-500/10 text-rose-400 border border-rose-500/20';
      case 'high':
      case 'urgency':
        return 'bg-orange-500/10 text-orange-400 border border-orange-500/20';
      case 'medium':
      case 'scarcity':
        return 'bg-amber-500/10 text-amber-400 border border-amber-500/20';
      case 'low':
      case 'social proof':
      case 'misdirection':
        return 'bg-indigo-500/10 text-indigo-400 border border-indigo-500/20';
      default:
        return 'bg-slate-500/15 text-slate-400 border border-white/5';
    }
  };

  const totalDatasetCount = Object.values(categoryCounts).reduce((a, b) => a + b, 0);
  const cleanCopyCount = categoryCounts['Not Dark Pattern'] || 0;
  const deceptiveClassCount = metrics ? Math.max(metrics.numClasses - 1, 0) : 0;

  return (
    <div className="space-y-6">
      {/* Model Performance Metrics Card */}
      <div className="grid md:grid-cols-3 gap-6">
        {/* Model Architecture */}
        <Card className="p-5 bg-card/40 backdrop-blur-md border border-border rounded-2xl flex flex-col justify-between">
          <div>
            <div className="flex items-center gap-2 mb-3">
              <Cpu className="w-5 h-5 text-indigo-400" />
              <h4 className="font-bold text-xs uppercase tracking-wider text-slate-300">Model Architecture</h4>
            </div>
            <div className="space-y-2 mt-2">
              <div className="flex justify-between text-xs">
                <span className="text-slate-500 font-medium">Pipeline:</span>
                <span className="text-slate-300 font-bold font-mono">TF-IDF + LogReg</span>
              </div>
              <div className="flex justify-between text-xs">
                <span className="text-slate-500 font-medium">Word / Character Features:</span>
                <span className="text-slate-300 font-bold font-mono">
                  {metrics
                    ? `${metrics.wordFeatureCount.toLocaleString()} / ${metrics.charFeatureCount.toLocaleString()}`
                    : '...'}
                </span>
              </div>
              <div className="flex justify-between text-xs">
                <span className="text-slate-500 font-medium">Word / Character N-grams:</span>
                <span className="text-slate-300 font-bold font-mono">
                  {metrics
                    ? `${metrics.ngramRange.join('-')} / ${metrics.charNgramRange.join('-')}`
                    : '...'}
                </span>
              </div>
            </div>
          </div>
          <div className="mt-4 pt-3 border-t border-white/5 text-[10px] text-slate-400 leading-relaxed flex items-center gap-1.5">
            <Settings className="w-3.5 h-3.5 text-indigo-400 animate-spin-slow" />
            Auto-retrained on Python backend startup.
          </div>
        </Card>

        {/* Model Statistics */}
        <Card className="p-5 bg-card/40 backdrop-blur-md border border-border rounded-2xl flex flex-col justify-between">
          <div>
            <div className="flex items-center gap-2 mb-3">
              <LineChart className="w-5 h-5 text-emerald-400" />
              <h4 className="font-bold text-xs uppercase tracking-wider text-slate-300">Model Performance</h4>
            </div>
            <div className="space-y-2 mt-2">
              <div className="flex justify-between text-xs">
                <span className="text-slate-500 font-medium">Classification Accuracy:</span>
                <span className="text-emerald-400 font-extrabold font-mono">
                  {metrics ? `${(metrics.accuracy * 100).toFixed(1)}%` : '...'}
                </span>
              </div>
              <div className="flex justify-between text-xs">
                <span className="text-slate-500 font-medium">F1 Score (Weighted):</span>
                <span className="text-emerald-400 font-extrabold font-mono">
                  {metrics ? `${(metrics.weightedF1 * 100).toFixed(1)}%` : '...'}
                </span>
              </div>
              <div className="flex justify-between text-xs">
                <span className="text-slate-500 font-medium">Precision / Recall:</span>
                <span className="text-emerald-400 font-extrabold font-mono">
                  {metrics
                    ? `${(metrics.weightedPrecision * 100).toFixed(1)}% / ${(metrics.weightedRecall * 100).toFixed(1)}%`
                    : '...'}
                </span>
              </div>
            </div>
          </div>
          <div className="mt-4 pt-3 border-t border-white/5 text-[10px] text-slate-400 leading-relaxed">
            {metrics
              ? `${metrics.splitMethod}; page overlap: ${metrics.groupOverlap}; confidence calibrated: ${metrics.calibrated ? 'yes' : 'no'}.`
              : 'Loading validation method...'}
          </div>
        </Card>

        {/* Dataset Stats */}
        <Card className="p-5 bg-card/40 backdrop-blur-md border border-border rounded-2xl flex flex-col justify-between">
          <div>
            <div className="flex items-center gap-2 mb-3">
              <Database className="w-5 h-5 text-cyan-400" />
              <h4 className="font-bold text-xs uppercase tracking-wider text-slate-300">Training Corpus</h4>
            </div>
            <div className="space-y-2 mt-2">
              <div className="flex justify-between text-xs">
                <span className="text-slate-500 font-medium">Total Training Examples:</span>
                <span className="text-white font-extrabold font-mono">
                  {metrics?.datasetSize || totalDatasetCount || '...'} rows
                </span>
              </div>
              <div className="flex justify-between text-xs">
                <span className="text-slate-500 font-medium">Deceptive Classes:</span>
                <span className="text-white font-extrabold font-mono">
                  {metrics ? deceptiveClassCount : '...'} Categories
                </span>
              </div>
              <div className="flex justify-between text-xs">
                <span className="text-slate-500 font-medium">Clean Interface Copy:</span>
                <span className="text-white font-extrabold font-mono">
                  {cleanCopyCount ? cleanCopyCount.toLocaleString() : '...'} rows
                </span>
              </div>
            </div>
          </div>
          <div className="mt-4 pt-3 border-t border-white/5 text-[10px] text-slate-400 leading-relaxed">
            Counts are loaded directly from the labeled NLP training corpus.
          </div>
        </Card>
      </div>

      {/* Dataset Chart Distributions */}
      <Card className="p-6 bg-card/40 backdrop-blur-md border border-border rounded-2xl">
        <h4 className="font-bold text-xs uppercase tracking-wider text-slate-300 mb-4 flex items-center gap-2">
          <Sparkles className="w-4 h-4 text-indigo-400 animate-pulse" />
          Pattern Category Distribution (Training Corpus)
        </h4>

        {isLoading ? (
          <div className="h-48 flex items-center justify-center text-xs text-slate-500">
            Loading distribution analytics...
          </div>
        ) : (
          <div className="h-56 w-full">
            <ResponsiveContainer width="100%" height="100%">
              <BarChart data={chartData} margin={{ top: 10, right: 10, left: -20, bottom: 5 }}>
                <XAxis dataKey="name" stroke="#64748b" fontSize={10} tickLine={false} />
                <YAxis stroke="#64748b" fontSize={10} tickLine={false} axisLine={false} />
                <Tooltip
                  contentStyle={{ backgroundColor: '#1e293b', border: '1px solid rgba(255,255,255,0.08)', borderRadius: '8px' }}
                  labelStyle={{ color: '#94a3b8', fontSize: '10px' }}
                  itemStyle={{ color: '#f8fafc', fontSize: '12px', fontWeight: 'bold' }}
                />
                <Bar dataKey="count" fill="#6366f1" radius={[4, 4, 0, 0]} maxBarSize={45} />
              </BarChart>
            </ResponsiveContainer>
          </div>
        )}
      </Card>

      {/* Searchable Training Grid */}
      <Card className="p-6 bg-card/40 backdrop-blur-md border border-border rounded-2xl">
        <div className="flex flex-col sm:flex-row sm:items-center justify-between gap-4 mb-6">
          <h4 className="font-bold text-xs uppercase tracking-wider text-slate-300">
            Training Records Browser
          </h4>

          {/* Filters */}
          <div className="flex items-center gap-2">
            {/* Category selection */}
            <select
              value={selectedCategory}
              onChange={(e) => handleCategoryChange(e.target.value)}
              className="bg-slate-950/40 text-slate-200 border border-border text-xs px-3 py-2 rounded-xl focus:outline-none focus:ring-1 focus:ring-indigo-500"
            >
              <option value="">All Categories</option>
              <option value="scarcity">Scarcity</option>
              <option value="urgency">Urgency</option>
              <option value="social proof">Social Proof</option>
              <option value="misdirection">Misdirection</option>
              <option value="sneaking">Sneaking</option>
              <option value="obstruction">Obstruction</option>
              <option value="forced action">Forced Action</option>
              <option value="not dark pattern">Not Dark Pattern</option>
            </select>

            {/* Search Input */}
            <div className="relative">
              <Search className="absolute left-3 top-2.5 w-4 h-4 text-slate-500" />
              <Input
                placeholder="Search training copy..."
                value={searchQuery}
                onChange={handleSearchChange}
                className="bg-slate-950/40 border-border text-white text-xs pl-9 pr-4 py-2 w-56 rounded-xl"
              />
            </div>
          </div>
        </div>

        {/* Data Table */}
        <div className="border border-border rounded-xl overflow-hidden bg-slate-950/20">
          <table className="w-full text-left border-collapse">
            <thead>
              <tr className="border-b border-border bg-slate-900/40 text-[10px] uppercase font-bold text-slate-400 tracking-wider">
                <th className="p-4 w-20">Row ID</th>
                <th className="p-4 w-40">Category</th>
                <th className="p-4">Training Text Copy</th>
                <th className="p-4 w-24 text-center">Label</th>
              </tr>
            </thead>
            <tbody className="divide-y divide-border text-xs">
              {isTableLoading ? (
                <tr>
                  <td colSpan={4} className="p-8 text-center text-slate-500 font-medium">
                    Querying model datasets...
                  </td>
                </tr>
              ) : records.length > 0 ? (
                records.map((record) => (
                  <tr key={record.page_id} className="hover:bg-white/5 transition-colors">
                    <td className="p-4 font-mono text-slate-500">#{record.page_id}</td>
                    <td className="p-4">
                      <Badge className={`${getSeverityBadgeClass(record['Pattern Category'])} text-[9px] font-bold px-2 py-0.5 rounded-md`}>
                        {record['Pattern Category']}
                      </Badge>
                    </td>
                    <td className="p-4 text-slate-200 leading-normal max-w-lg break-words">
                      {record.text}
                    </td>
                    <td className="p-4 font-mono text-slate-400 text-center font-bold">
                      {record.label}
                    </td>
                  </tr>
                ))
              ) : (
                <tr>
                  <td colSpan={4} className="p-8 text-center text-slate-500">
                    No matching training samples found in dataset.
                  </td>
                </tr>
              )}
            </tbody>
          </table>
        </div>

        {/* Pagination controls */}
        {totalRecords > PAGE_LIMIT && (
          <div className="flex items-center justify-between mt-5 border-t border-white/5 pt-4">
            <span className="text-[10px] font-bold text-slate-500 uppercase tracking-wider">
              Showing {currentPage * PAGE_LIMIT + 1} - {Math.min((currentPage + 1) * PAGE_LIMIT, totalRecords)} of {totalRecords} records
            </span>

            <div className="flex items-center gap-2">
              <Button
                variant="outline"
                size="sm"
                onClick={() => setCurrentPage((p) => Math.max(0, p - 1))}
                disabled={currentPage === 0 || isTableLoading}
                className="h-8 border-border text-slate-300 hover:text-white rounded-xl px-2"
              >
                <ChevronLeft className="w-4 h-4" />
              </Button>
              <span className="text-xs font-bold text-slate-300 bg-slate-900/60 border border-border px-3 py-1.5 rounded-xl font-mono">
                {currentPage + 1} / {Math.ceil(totalRecords / PAGE_LIMIT)}
              </span>
              <Button
                variant="outline"
                size="sm"
                onClick={() => setCurrentPage((p) => Math.min(Math.ceil(totalRecords / PAGE_LIMIT) - 1, p + 1))}
                disabled={(currentPage + 1) * PAGE_LIMIT >= totalRecords || isTableLoading}
                className="h-8 border-border text-slate-300 hover:text-white rounded-xl px-2"
              >
                <ChevronRight className="w-4 h-4" />
              </Button>
            </div>
          </div>
        )}
      </Card>
    </div>
  );
}