App.tsx

import React, { useState, useEffect } from 'react';
import JSZip from 'jszip';
import { Upload, Play, Settings, CheckCircle, AlertTriangle, UserX, FileImage, Download, BrainCircuit, X, ImagePlus, ChevronRight, Loader, FolderPlus, Trash2 } from 'lucide-react';
import { GoogleGenAI } from "@google/genai";

// Local imports (Mocking types locally for clarity since structure changed)
// In your real project, update types.ts to match these interfaces if needed.
import { AppState as BaseAppState, ImageFile, ProcessingConfig } from './types';
import { initPoseLandmarker, detectPose, normalizeLandmarks, calculatePoseDistance } from './services/poseService';

// --- New Interfaces for Multi-Reference Support ---
interface ReferenceImage {
  id: string;
  previewUrl: string;
  file: File;
  landmarks?: any; // NormalizedLandmarkList
}

interface ReferenceCategory {
  id: string;
  name: string;
  images: ReferenceImage[];
}

interface AppState extends Omit<BaseAppState, 'references'> {
  categories: ReferenceCategory[]; // Replaces flat 'references'
}

const DEFAULT_CONFIG: ProcessingConfig = {
  confidenceThreshold: 0.5,
  maxDistanceThreshold: 0.45,
};

const App: React.FC = () => {
  const [state, setState] = useState<AppState>({
    step: 'setup',
    config: DEFAULT_CONFIG,
    categories: [], // New structure
    dataset: [],
    isProcessing: false,
    progress: 0,
  });

  const [analyzingId, setAnalyzingId] = useState<string | null>(null);
  const [isZipping, setIsZipping] = useState(false);

  useEffect(() => {
    initPoseLandmarker().catch(console.error);
  }, []);

  // --- Handlers ---

  const handleConfigChange = (key: keyof ProcessingConfig, value: number) => {
    setState(prev => ({ ...prev, config: { ...prev.config, [key]: value } }));
  };

  // 1. Create a new Category
  const handleAddCategory = () => {
    const newCat: ReferenceCategory = {
      id: Math.random().toString(36).substr(2, 9),
      name: `Step ${state.categories.length + 1}`,
      images: []
    };
    setState(prev => ({ ...prev, categories: [...prev.categories, newCat] }));
  };

  // 2. Remove a Category
  const handleRemoveCategory = (catId: string) => {
    setState(prev => ({ ...prev, categories: prev.categories.filter(c => c.id !== catId) }));
  };

  // 3. Rename Category
  const handleUpdateCategoryName = (catId: string, name: string) => {
    setState(prev => ({
      ...prev,
      categories: prev.categories.map(c => c.id === catId ? { ...c, name } : c)
    }));
  };

  // 4. Upload images to a specific Category
  const handleReferenceUpload = (e: React.ChangeEvent<HTMLInputElement>, catId: string) => {
    if (e.target.files) {
      const newRefs: ReferenceImage[] = [];
      const files = Array.from(e.target.files) as File[];
      for (const file of files) {
        newRefs.push({
          id: Math.random().toString(36).substr(2, 9),
          file,
          previewUrl: URL.createObjectURL(file),
        });
      }
      setState(prev => ({
        ...prev,
        categories: prev.categories.map(c => 
          c.id === catId ? { ...c, images: [...c.images, ...newRefs] } : c
        )
      }));
    }
  };

  // 5. Remove image from category
  const handleRemoveRefImage = (catId: string, imgId: string) => {
    setState(prev => ({
      ...prev,
      categories: prev.categories.map(c => 
        c.id === catId ? { ...c, images: c.images.filter(img => img.id !== imgId) } : c
      )
    }));
  };

  const handleDatasetUpload = (e: React.ChangeEvent<HTMLInputElement>) => {
    if (e.target.files) {
      const newFiles: ImageFile[] = (Array.from(e.target.files) as File[]).map(file => ({
        id: Math.random().toString(36).substr(2, 9),
        file,
        previewUrl: URL.createObjectURL(file),
        status: 'unprocessed'
      }));
      setState(prev => ({ ...prev, dataset: [...prev.dataset, ...newFiles] }));
    }
  };

  // --- Main Logic ---

  const runClassification = async () => {
    const totalRefs = state.categories.reduce((acc, cat) => acc + cat.images.length, 0);
    if (totalRefs === 0) {
      alert("Please add at least one reference image in a category.");
      return;
    }

    setState(prev => ({ ...prev, step: 'processing', isProcessing: true, progress: 0 }));

    // 1. Pre-process Reference Images (Calculate Landmarks for ALL refs in ALL categories)
    const processedCategories = await Promise.all(state.categories.map(async (cat) => {
      const processedImages = await Promise.all(cat.images.map(async (ref) => {
        const landmarks = await detectPose(ref.previewUrl);
        if (!landmarks) return { ...ref, landmarks: undefined };
        const normalized = normalizeLandmarks(landmarks);
        return { ...ref, landmarks: normalized || undefined };
      }));
      
      // Filter out invalid images inside the category
      const validImages = processedImages.filter(img => img.landmarks !== undefined);
      return { ...cat, images: validImages };
    }));

    // Filter out empty categories
    const validCategories = processedCategories.filter(c => c.images.length > 0);

    if (validCategories.length === 0) {
      alert("No valid poses detected in any reference images!");
      setState(prev => ({ ...prev, isProcessing: false, step: 'setup' }));
      return;
    }

    // 2. Process Dataset
    const dataset = [...state.dataset];
    let completed = 0;
    const BATCH_SIZE = 5;

    for (let i = 0; i < dataset.length; i += BATCH_SIZE) {
      const batch = dataset.slice(i, i + BATCH_SIZE);
      
      await Promise.all(batch.map(async (img) => {
        const rawLandmarks = await detectPose(img.previewUrl);

        if (!rawLandmarks) {
          img.status = 'no_person';
          return;
        }

        const avgVis = rawLandmarks.reduce((acc, l) => acc + (l.visibility || 0), 0) / rawLandmarks.length;
        if (avgVis < state.config.confidenceThreshold) {
          img.status = 'no_person';
          return;
        }

        const normalizedInput = normalizeLandmarks(rawLandmarks);
        if (!normalizedInput) {
          img.status = 'no_person';
          return;
        }

        // --- Multi-Reference Matching Logic ---
        let globalMinDistance = Infinity;
        let bestMatchCategoryId = '';

        // Iterate over each Category
        for (const category of validCategories) {
          // Find the BEST match (closest distance) within this category's images
          let categoryMinDist = Infinity;
          
          for (const refImg of category.images) {
             const dist = calculatePoseDistance(normalizedInput, refImg.landmarks);
             if (dist < categoryMinDist) {
               categoryMinDist = dist;
             }
          }

          // If this category's best match is better than the global best match
          if (categoryMinDist < globalMinDistance) {
            globalMinDistance = categoryMinDist;
            bestMatchCategoryId = category.id;
          }
        }

        img.distance = globalMinDistance;

        if (globalMinDistance > state.config.maxDistanceThreshold) {
          img.status = 'uncertain';
        } else {
          img.status = 'classified';
          img.assignedCategory = bestMatchCategoryId;
        }
      }));

      completed += batch.length;
      setState(prev => ({ 
        ...prev, 
        progress: (completed / dataset.length) * 100 
      }));
      
      await new Promise(r => setTimeout(r, 0));
    }

    setState(prev => ({ 
      ...prev, 
      dataset: dataset,
      isProcessing: false, 
      step: 'results',
      categories: validCategories // Update with landmarks
    }));
  };

  // --- Gemini & Download Helpers ---

  const analyzeWithGemini = async (image: ImageFile) => {
     if (!process.env.API_KEY) {
       alert("API Key not found in environment.");
       return;
     }
     setAnalyzingId(image.id);
     try {
       const ai = new GoogleGenAI({ apiKey: process.env.API_KEY });
       const response = await fetch(image.previewUrl);
       const blob = await response.blob();
       const reader = new FileReader();
       const base64Promise = new Promise<string>((resolve) => {
         reader.onloadend = () => resolve((reader.result as string).split(',')[1]);
         reader.readAsDataURL(blob);
       });
       const base64Data = await base64Promise;
       const modelResponse = await ai.models.generateContent({
         model: 'gemini-2.5-flash',
         contents: {
           parts: [
             { inlineData: { mimeType: blob.type, data: base64Data } },
             { text: "Analyze this factory worker image. Is there a person? What action? Brief (1 sentence)." }
           ]
         }
       });
       setState(prev => ({
         ...prev,
         dataset: prev.dataset.map(img => img.id === image.id ? { ...img, geminiAnalysis: modelResponse.text } : img)
       }));
     } catch (error) {
       console.error("Gemini Error:", error);
       alert("Failed to analyze.");
     } finally {
       setAnalyzingId(null);
     }
  };

  const analyzeAllUncertainWithGemini = async () => {
    const uncertainImages = state.dataset.filter(img => img.status === 'uncertain' && !img.geminiAnalysis);
    if (uncertainImages.length === 0) { alert("No new uncertain images."); return; }
    if (!process.env.API_KEY) { alert("API Key missing."); return; }
    if (!confirm(`Analyze ${uncertainImages.length} images with Gemini?`)) return;
    for (const image of uncertainImages) await analyzeWithGemini(image);
    alert("Analysis complete.");
  };

  const downloadZip = async () => {
    if (isZipping) return;
    setIsZipping(true);
    console.log("Starting ZIP generation...");
    
    try {
      const zip = new JSZip();
      const uncertainFolder = zip.folder("_Uncertain_Needs_Review");
      const noPersonFolder = zip.folder("_No_Person_Detected");
      
      const categoryFolders: Record<string, JSZip | null> = {};
      state.categories.forEach(cat => {
        categoryFolders[cat.id] = zip.folder(cat.name);
      });

      for (const img of state.dataset) {
        let folder;
        if (img.status === 'no_person') folder = noPersonFolder;
        else if (img.status === 'uncertain') folder = uncertainFolder;
        else if (img.status === 'classified' && img.assignedCategory) {
          folder = categoryFolders[img.assignedCategory];
        }
        if(folder && img.file) folder.file(img.file.name, img.file);
      }

      const content = await zip.generateAsync({ type: "blob" });
      const url = URL.createObjectURL(content);
      const link = document.createElement("a");
      link.href = url;
      link.download = "classified_assembly_dataset.zip";
      document.body.appendChild(link);
      link.click();
      
      setTimeout(() => {
        document.body.removeChild(link);
        URL.revokeObjectURL(url);
      }, 2000);

    } catch (error) {
      console.error(error);
      alert("ZIP creation failed.");
    } finally {
      setIsZipping(false);
    }
  };

  const downloadSelectiveZip = async (zipName: string, images: ImageFile[]) => {
    if (images.length === 0 || isZipping) return;
    setIsZipping(true);
    try {
      const zip = new JSZip();
      for (const img of images) if(img.file) zip.file(img.file.name, img.file);
      const content = await zip.generateAsync({ type: "blob" });
      const url = URL.createObjectURL(content);
      const link = document.createElement("a");
      link.href = url;
      link.download = `${zipName.replace(/[^a-z0-9]/gi, '_')}.zip`;
      document.body.appendChild(link);
      link.click();
      setTimeout(() => { document.body.removeChild(link); URL.revokeObjectURL(url); }, 2000);
    } catch (error) {
      console.error(error);
      alert("ZIP creation failed.");
    } finally {
      setIsZipping(false);
    }
  };

  // --- Render Components ---

  const renderSetup = () => (
    <div className="space-y-6 animate-fade-in">
      <div className="bg-slate-800 p-6 rounded-lg border border-slate-700">
        <h2 className="text-xl font-semibold mb-4 flex items-center gap-2">
          <Settings className="text-blue-400" />
          1. Define Categories & Upload References
        </h2>
        <p className="text-slate-400 mb-4 text-sm">
          Create categories for each assembly step. Add 3-5 reference images per category for better accuracy.
        </p>
        
        {/* Category List */}
        <div className="space-y-6">
          {state.categories.map((cat, index) => (
            <div key={cat.id} className="bg-slate-900/50 border border-slate-600 rounded-lg p-4">
              <div className="flex items-center justify-between mb-3">
                <div className="flex items-center gap-3 flex-1">
                  <span className="bg-slate-700 text-xs px-2 py-1 rounded text-slate-300">Class {index + 1}</span>
                  <input 
                    type="text" 
                    value={cat.name}
                    onChange={(e) => handleUpdateCategoryName(cat.id, e.target.value)}
                    className="bg-slate-800 border-b border-slate-500 focus:border-blue-500 outline-none text-white font-medium px-2 py-1 w-full max-w-xs"
                    placeholder="Category Name (e.g. Tightening Screws)"
                  />
                </div>
                <button 
                  onClick={() => handleRemoveCategory(cat.id)}
                  className="text-slate-500 hover:text-red-400 p-2"
                  title="Remove Category"
                >
                  <Trash2 size={18} />
                </button>
              </div>

              {/* Images Grid for this Category */}
              <div className="grid grid-cols-2 md:grid-cols-6 gap-3">
                {cat.images.map(img => (
                  <div key={img.id} className="relative group h-24 w-full">
                    <img src={img.previewUrl} className="w-full h-full object-cover rounded border border-slate-700" />
                    <button 
                      onClick={() => handleRemoveRefImage(cat.id, img.id)}
                      className="absolute top-1 right-1 bg-red-500/80 hover:bg-red-600 text-white p-1 rounded-full opacity-0 group-hover:opacity-100 transition-opacity"
                    >
                      <X size={10} />
                    </button>
                  </div>
                ))}
                
                {/* Upload Button specific to this category */}
                <label className="border-2 border-dashed border-slate-700 hover:border-blue-500 hover:bg-slate-800 rounded flex flex-col items-center justify-center h-24 cursor-pointer transition-colors">
                  <ImagePlus size={20} className="text-slate-500 mb-1" />
                  <span className="text-[10px] text-slate-500">Add Images</span>
                  <input 
                    type="file" 
                    multiple 
                    accept="image/*" 
                    onChange={(e) => handleReferenceUpload(e, cat.id)} 
                    className="hidden" 
                  />
                </label>
              </div>
            </div>
          ))}
        </div>

        {/* Add Category Button */}
        <button 
          onClick={handleAddCategory}
          className="w-full mt-4 py-3 border-2 border-dashed border-slate-600 text-slate-400 hover:text-white hover:border-blue-500 hover:bg-slate-800/50 rounded-lg flex items-center justify-center gap-2 transition-all"
        >
          <FolderPlus size={20} />
          Create New Category
        </button>

      </div>

      <div className="flex justify-end">
        <button 
          onClick={() => setState(prev => ({ ...prev, step: 'upload' }))}
          disabled={state.categories.length === 0 || state.categories.some(c => c.images.length === 0)}
          className="bg-blue-600 hover:bg-blue-500 disabled:opacity-50 disabled:cursor-not-allowed px-6 py-2 rounded-lg font-medium flex items-center gap-2"
        >
          Next: Upload Dataset <ChevronRight size={16} />
        </button>
      </div>
    </div>
  );

  const renderUpload = () => (
    <div className="space-y-6 animate-fade-in">
      <div className="bg-slate-800 p-6 rounded-lg border border-slate-700">
        <h2 className="text-xl font-semibold mb-4 flex items-center gap-2">
          <Upload className="text-blue-400" />
          2. Upload Unclassified Dataset
        </h2>
        
        <div className="border-2 border-dashed border-slate-600 rounded-xl p-10 flex flex-col items-center justify-center bg-slate-900/50">
           <input type="file" multiple accept="image/*" onChange={handleDatasetUpload} className="hidden" id="dataset-upload" />
          <label htmlFor="dataset-upload" className="cursor-pointer flex flex-col items-center">
            <FileImage size={48} className="text-slate-500 mb-4" />
            <span className="text-lg font-medium text-slate-300">
              {state.dataset.length > 0 ? `${state.dataset.length} images loaded` : "Click to select images"}
            </span>
          </label>
        </div>
        
        {state.dataset.length > 0 && (
           <div className="mt-6 p-4 bg-slate-900 rounded border border-slate-700">
             <h3 className="text-sm font-semibold text-slate-300 mb-3">Classification Parameters</h3>
             <div className="grid grid-cols-1 md:grid-cols-2 gap-6">
               <div>
                 <div className="flex justify-between text-xs mb-1">
                   <span className="text-slate-400">MediaPipe Confidence</span>
                   <span className="text-blue-400">{state.config.confidenceThreshold.toFixed(2)}</span>
                 </div>
                 <input type="range" min="0" max="1" step="0.05" value={state.config.confidenceThreshold}
                  onChange={(e) => handleConfigChange('confidenceThreshold', parseFloat(e.target.value))}
                  className="w-full h-2 bg-slate-700 rounded-lg cursor-pointer accent-blue-500" />
               </div>
               <div>
                 <div className="flex justify-between text-xs mb-1">
                   <span className="text-slate-400">Strictness</span>
                   <span className="text-blue-400">{state.config.maxDistanceThreshold.toFixed(2)}</span>
                 </div>
                 <input type="range" min="0.1" max="1.5" step="0.05" value={state.config.maxDistanceThreshold}
                  onChange={(e) => handleConfigChange('maxDistanceThreshold', parseFloat(e.target.value))}
                  className="w-full h-2 bg-slate-700 rounded-lg cursor-pointer accent-blue-500" />
               </div>
             </div>
           </div>
        )}
      </div>

      <div className="flex justify-between">
         <button onClick={() => setState(prev => ({ ...prev, step: 'setup' }))} className="text-slate-400 hover:text-white px-4 py-2">Back</button>
        <button onClick={runClassification} disabled={state.dataset.length === 0}
          className="bg-green-600 hover:bg-green-500 disabled:opacity-50 px-6 py-2 rounded-lg font-medium flex items-center gap-2 shadow-lg">
          <Play size={16} fill="currentColor" /> Start Classification
        </button>
      </div>
    </div>
  );

  const renderResults = () => {
    // Group results logic updated for Categories
    const grouped: Record<string, ImageFile[]> = { uncertain: [], no_person: [] };
    state.categories.forEach(c => grouped[c.id] = []); // Init category arrays

    state.dataset.forEach(img => {
      if (img.status === 'no_person') grouped.no_person.push(img);
      else if (img.status === 'uncertain') grouped.uncertain.push(img);
      else if (img.assignedCategory && grouped[img.assignedCategory]) grouped[img.assignedCategory].push(img);
    });

    return (
      <div className="space-y-8 animate-fade-in pb-20">
        <div className="flex justify-between items-center bg-slate-800 p-4 rounded-lg border border-slate-700 sticky top-4 z-10 shadow-xl">
          <div>
            <h2 className="text-lg font-bold">Classification Complete</h2>
            <p className="text-xs text-slate-400">Processed {state.dataset.length} images</p>
          </div>
          <div className="flex gap-3">
             <button onClick={() => setState(prev => ({ ...prev, step: 'upload' }))} className="px-4 py-2 text-sm text-slate-300 hover:text-white">Restart</button>
            <button onClick={analyzeAllUncertainWithGemini} className="bg-purple-600 hover:bg-purple-500 px-4 py-2 rounded flex items-center gap-2 text-sm font-medium" disabled={grouped.uncertain.length === 0}>
              <BrainCircuit size={16} /> AI Analyze
            </button>
            <button onClick={downloadZip} disabled={isZipping} className="bg-blue-600 hover:bg-blue-500 disabled:bg-slate-600 px-4 py-2 rounded flex items-center gap-2 text-sm font-medium shadow-lg">
              {isZipping ? <Loader size={16} className="animate-spin" /> : <Download size={16} />}
              {isZipping ? "Compressing..." : "Download ZIP"}
            </button>
          </div>
        </div>

        {/* Uncertain Section */}
        <section className="bg-slate-800/50 border border-yellow-700/30 rounded-lg p-6">
           <div className="flex justify-between items-center mb-4">
            <div className="flex items-center gap-2 text-yellow-500">
              <AlertTriangle />
              <h3 className="text-lg font-semibold">_Uncertain ({grouped.uncertain.length})</h3>
            </div>
            <button onClick={() => downloadSelectiveZip('_Uncertain', grouped.uncertain)} disabled={grouped.uncertain.length === 0 || isZipping} className="bg-yellow-700 hover:bg-yellow-600 disabled:bg-slate-600 px-3 py-1 rounded flex items-center gap-2 text-xs font-medium text-yellow-100">
              <Download size={14} /> Download
            </button>
          </div>
          <div className="grid grid-cols-2 md:grid-cols-5 gap-4">
            {grouped.uncertain.slice(0, 10).map(img => (
              <div key={img.id} className="bg-slate-900 rounded border border-slate-700 p-2 group relative">
                <img src={img.previewUrl} className="w-full h-32 object-cover rounded opacity-70" />
                <span className="absolute bottom-1 right-1 text-[10px] bg-black/50 text-white px-1 rounded">D: {img.distance?.toFixed(3)}</span>
                 <button onClick={() => analyzeWithGemini(img)} disabled={analyzingId === img.id} className="absolute top-1 right-1 bg-purple-900/80 text-white px-2 py-1 rounded text-[10px] opacity-0 group-hover:opacity-100 transition-opacity">AI Analyze</button>
              </div>
            ))}
          </div>
        </section>

        {/* Categories Section */}
        <div className="grid grid-cols-1 gap-6">
          {state.categories.map(cat => {
              const images = grouped[cat.id] || [];
              return (
                <section key={cat.id} className="bg-slate-800 rounded-lg p-6 border border-slate-700">
                   <div className="flex justify-between items-center mb-4">
                    <div className="flex items-center gap-2 text-green-400">
                      <CheckCircle size={20} />
                      <h3 className="text-lg font-semibold">{cat.name} ({images.length})</h3>
                    </div>
                    <button onClick={() => downloadSelectiveZip(cat.name, images)} disabled={images.length === 0 || isZipping} className="bg-blue-700 hover:bg-blue-600 disabled:bg-slate-600 px-3 py-1 rounded flex items-center gap-2 text-xs font-medium">
                      <Download size={14} /> Download
                    </button>
                  </div>
                  <div className="grid grid-cols-3 md:grid-cols-8 gap-2">
                    {images.slice(0, 8).map(img => (
                       <div key={img.id} className="relative group">
                         <img src={img.previewUrl} className="w-full h-24 object-cover rounded border border-slate-600" />
                         <div className="absolute inset-0 bg-black/50 opacity-0 group-hover:opacity-100 flex items-center justify-center text-[10px] text-white rounded pointer-events-none">
                           D: {img.distance?.toFixed(3)}
                         </div>
                       </div>
                    ))}
                     {images.length > 8 && <div className="flex items-center justify-center text-slate-500 text-xs bg-slate-900 rounded h-24">+{images.length - 8}</div>}
                  </div>
                </section>
              )
          })}
        </div>
      </div>
    );
  };

  return (
    <div className="min-h-screen bg-slate-950 text-slate-200 font-sans selection:bg-blue-500/30">
      <header className="bg-slate-900 border-b border-slate-800 py-4 px-6 flex justify-between items-center">
        <div className="flex items-center gap-3">
          <div className="bg-blue-600 p-2 rounded-lg"><BrainCircuit size={24} className="text-white" /></div>
          <div><h1 className="text-xl font-bold tracking-tight text-white">AssemblyPose Pro <span className="text-xs font-normal text-yellow-400 border border-yellow-600/50 px-1 rounded ml-2">Multi-Ref</span></h1></div>
        </div>
        <div className="text-xs text-slate-600 hidden md:block">MediaPipe Vision + Gemini 2.5</div>
      </header>

      <main className="max-w-6xl mx-auto p-6">
        {state.step === 'setup' && renderSetup()}
        {state.step === 'upload' && renderUpload()}
        {state.step === 'processing' && (
           <div className="flex flex-col items-center justify-center h-64 animate-fade-in">
            <div className="w-full max-w-md bg-slate-800 rounded-full h-4 mb-4 overflow-hidden border border-slate-700">
              <div className="bg-blue-500 h-full transition-all duration-300 ease-out" style={{ width: `${state.progress}%` }} />
            </div>
            <p className="text-slate-300 font-mono text-lg">{Math.round(state.progress)}%</p>
          </div>
        )}
        {state.step === 'results' && renderResults()}
      </main>
    </div>
  );
};

export default App;