src/components/viewer-tab.tsx

"use client";

import q3Results from "../../results/qwen3.json"
import q3_30B_A3B_Results from "../../results/qwen3-30B-A3-results.json"
// import mockResults from "../../qwen3-final-results.json"
import { useMemo, useState, useEffect, useRef } from "react";
import { Card } from "@/components/ui/card";
import ForceDirectedGraph from "@/components/force-directed-graph";
import RunsList from "@/components/runs-list";
import { 
  Select, 
  SelectContent, 
  SelectItem, 
  SelectTrigger, 
  SelectValue 
} from "@/components/ui/select";
import { Run as ForceGraphRun } from "@/components/reasoning-trace";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { UploadIcon } from "lucide-react";

const defaultModels = {
  "Qwen3-14B": q3Results,
  "Qwen3-30B-A3B": q3_30B_A3B_Results,
}

// Use the type expected by RunsList
interface Run {
  start_article: string;
  destination_article: string;
  steps: string[];
  result: string;
}

// Interface for model statistics
interface ModelStats {
  winPercentage: number;
  avgSteps: number;
  stdDevSteps: number;
  totalRuns: number;
  wins: number;
  medianSteps: number;
  minSteps: number;
  maxSteps: number;
}

export default function ViewerTab({
  handleTryRun,
}: {
  handleTryRun: (startArticle: string, destinationArticle: string) => void;
}) {
  const [selectedRun, setSelectedRun] = useState<number | null>(null);
  const [runs, setRuns] = useState<Run[]>([]);
  const [selectedModel, setSelectedModel] = useState<string>("Qwen3-14B");
  const [modelStats, setModelStats] = useState<ModelStats | null>(null);
  const [models, setModels] = useState(defaultModels);
  const fileInputRef = useRef<HTMLInputElement>(null);

  useEffect(() => {
    // Convert the model data to the format expected by RunsList
    const convertedRuns = models[selectedModel]?.runs?.map((run: {
      start_article: string;
      destination_article: string;
      steps: { type: string; article: string }[];
      result: string;
    }) => ({
      start_article: run.start_article,
      destination_article: run.destination_article,
      steps: run.steps.map((step: { article: string }) => step.article),
      result: run.result
    })) || [];
    setRuns(convertedRuns);

    // Calculate model statistics
    const winRuns = convertedRuns.filter(run => run.result === "win");
    const totalRuns = convertedRuns.length;
    const wins = winRuns.length;
    const winPercentage = totalRuns > 0 ? (wins / totalRuns) * 100 : 0;
    
    // Calculate steps statistics for winning runs
    const stepCounts = winRuns.map(run => run.steps.length);
    const avgSteps = stepCounts.length > 0 
      ? stepCounts.reduce((sum, count) => sum + count, 0) / stepCounts.length 
      : 0;
    
    // Calculate standard deviation
    const variance = stepCounts.length > 0
      ? stepCounts.reduce((sum, count) => sum + Math.pow(count - avgSteps, 2), 0) / stepCounts.length
      : 0;
    const stdDevSteps = Math.sqrt(variance);

    // Calculate median, min, max steps
    const sortedSteps = [...stepCounts].sort((a, b) => a - b);
    const medianSteps = stepCounts.length > 0
      ? stepCounts.length % 2 === 0
        ? (sortedSteps[stepCounts.length / 2 - 1] + sortedSteps[stepCounts.length / 2]) / 2
        : sortedSteps[Math.floor(stepCounts.length / 2)]
      : 0;
    const minSteps = stepCounts.length > 0 ? Math.min(...stepCounts) : 0;
    const maxSteps = stepCounts.length > 0 ? Math.max(...stepCounts) : 0;

    setModelStats({
      winPercentage,
      avgSteps,
      stdDevSteps,
      totalRuns,
      wins,
      medianSteps,
      minSteps,
      maxSteps
    });
  }, [selectedModel, models]);

  const handleRunSelect = (runId: number) => {
    setSelectedRun(runId);
  };

  const filterRuns = useMemo(() => {
    return runs.filter(run => run.result === "win");
  }, [runs]);

  // Convert the runs to the format expected by ForceDirectedGraph
  const forceGraphRuns = useMemo(() => {
    return filterRuns.map((run): ForceGraphRun => ({
      start_article: run.start_article,
      destination_article: run.destination_article,
      steps: run.steps.map(article => ({ type: "move", article }))
    }));
  }, [filterRuns]);

  const handleFileUpload = (event: React.ChangeEvent<HTMLInputElement>) => {
    const file = event.target.files?.[0];
    if (!file) return;

    const reader = new FileReader();
    reader.onload = (e) => {
      try {
        const jsonData = JSON.parse(e.target?.result as string);
        
        // Validate the JSON structure has the required fields
        if (!jsonData.runs || !Array.isArray(jsonData.runs)) {
          alert("Invalid JSON format. File must contain a 'runs' array.");
          return;
        }
        
        // Create a filename-based model name, removing extension and path
        const fileName = file.name.replace(/\.[^/.]+$/, "");
        const modelName = `Custom: ${fileName}`;
        
        // Add the new model to the models object
        setModels(prev => ({
          ...prev,
          [modelName]: jsonData
        }));
        
        // Select the newly added model
        setSelectedModel(modelName);
      } catch (error) {
        alert(`Error parsing JSON file: ${error.message}`);
      }
    };
    reader.readAsText(file);
    
    // Reset the file input
    if (fileInputRef.current) {
      fileInputRef.current.value = '';
    }
  };

  const handleUploadClick = () => {
    fileInputRef.current?.click();
  };

  return (
    <div className="grid grid-cols-1 md:grid-cols-12 gap-4 h-[calc(100vh-200px)] max-h-[calc(100vh-200px)] overflow-hidden p-2">
     <Card className="p-3 col-span-12 row-start-1">
       <div className="flex flex-col sm:flex-row items-start sm:items-center gap-3">
         <div className="flex-shrink-0">
           <Select value={selectedModel} onValueChange={setSelectedModel}>
             <SelectTrigger className="w-[180px]">
               <SelectValue placeholder="Select model" />
             </SelectTrigger>
             <SelectContent>
               {Object.keys(models).map((modelName) => (
                 <SelectItem key={modelName} value={modelName}>
                   {modelName}
                 </SelectItem>
               ))}
             </SelectContent>
           </Select>
         </div>
         
         <Button 
           variant="outline" 
           size="sm" 
           className="flex items-center gap-1" 
           onClick={handleUploadClick}
         >
           <UploadIcon size={14} />
           <span>Upload JSON</span>
           <input 
             type="file" 
             ref={fileInputRef}
             accept=".json" 
             className="hidden" 
             onChange={handleFileUpload} 
           />
         </Button>
         
         {modelStats && (
           <div className="flex flex-wrap gap-1.5 items-center">
             <Badge variant="outline" className="px-2 py-0.5 flex gap-1 items-center">
               <span className="text-xs font-medium">Success:</span>
               <span className="text-xs font-semibold">{modelStats.winPercentage.toFixed(1)}%</span>
               <span className="text-xs text-muted-foreground">({modelStats.wins}/{modelStats.totalRuns})</span>
             </Badge>
             
             <Badge variant="outline" className="px-2 py-0.5 flex gap-1 items-center">
               <span className="text-xs font-medium">Mean:</span>
               <span className="text-xs font-semibold">{modelStats.avgSteps.toFixed(1)}</span>
               <span className="text-xs text-muted-foreground">±{modelStats.stdDevSteps.toFixed(1)}</span>
             </Badge>
             
             <Badge variant="outline" className="px-2 py-0.5 flex gap-1 items-center">
               <span className="text-xs font-medium">Median:</span>
               <span className="text-xs font-semibold">{modelStats.medianSteps.toFixed(1)}</span>
             </Badge>
             
             <Badge variant="outline" className="px-2 py-0.5 flex gap-1 items-center">
               <span className="text-xs font-medium">Min:</span>
               <span className="text-xs font-semibold">{modelStats.minSteps}</span>
             </Badge>
             
             <Badge variant="outline" className="px-2 py-0.5 flex gap-1 items-center">
               <span className="text-xs font-medium">Max:</span>
               <span className="text-xs font-semibold">{modelStats.maxSteps}</span>
             </Badge>
           </div>
         )}
       </div>
     </Card>
      <div className="md:col-span-3 flex flex-col max-h-full overflow-hidden">
        <div className="bg-card rounded-lg p-3 border flex-grow overflow-hidden flex flex-col">
          <h3 className="text-sm font-medium mb-2 text-muted-foreground flex-shrink-0">
            Runs
          </h3>
          <div className="flex-grow overflow-hidden">
            <RunsList
              runs={filterRuns}
              onSelectRun={handleRunSelect}
              selectedRunId={selectedRun}
              onTryRun={handleTryRun}
            />
          </div>
        </div>
      </div>

      <div className="md:col-span-9 max-h-full overflow-hidden">
        <Card className="w-full h-full flex items-center justify-center p-0 m-0 overflow-hidden">
          <ForceDirectedGraph runs={forceGraphRuns} runId={selectedRun} />
        </Card>
      </div>
    </div>
  );
}