src/App.jsx

import { useState, useEffect, useRef } from "react";
import * as pdfjsLib from 'pdfjs-dist';
import pdfWorkerSrc from 'pdfjs-dist/build/pdf.worker.mjs?url';
pdfjsLib.GlobalWorkerOptions.workerSrc = pdfWorkerSrc;
import "./App.css";

function App() {
  // Worker reference
  const workerRef = useRef(null);

  // Application State
  const [modelStatus, setModelStatus] = useState("idle");
  const [progressItems, setProgressItems] = useState({});
  const [totalProgress, setTotalProgress] = useState(0);
  const [errorMsg, setErrorMsg] = useState("");
  const [loadTime, setLoadTime] = useState(null);

  // Performance metrics
  const [loadedDevice, setLoadedDevice] = useState("Detecting...");
  const [deviceConfig, setDeviceConfig] = useState("wasm");
  const [generationTime, setGenerationTime] = useState(null);

  // VLM State
  const [imageFile, setImageFile] = useState(null);
  const [imagePreviewUrl, setImagePreviewUrl] = useState("");
  const [extractedMetadata, setExtractedMetadata] = useState(null);
  const [isGenerating, setIsGenerating] = useState(false);

  // Initialize Web Worker
  useEffect(() => {
    workerRef.current = new Worker(new URL("./ai/worker.js", import.meta.url), {
      type: "module",
    });

    const onMessageReceived = (e) => {
      const { type, data, error } = e.data;

      switch (type) {
        case "progress":
          setProgressItems((prev) => {
            const newItems = { ...prev, [data.file]: data };
            const itemsArray = Object.values(newItems);
            const loaded = itemsArray.reduce((acc, item) => acc + (item.loaded || 0), 0);
            const total = itemsArray.reduce((acc, item) => acc + (item.total || 0), 0);
            if (total > 0) {
              setTotalProgress((loaded / total) * 100);
            }
            return newItems;
          });
          break;

        case "ready":
          setModelStatus("ready");
          setLoadedDevice(data.device === "webgpu" ? "WebGPU" : "WASM CPU");
          setLoadTime((performance.now() - loadTime) / 1000);
          break;

        case "completed":
          setExtractedMetadata(data.metadata);
          setGenerationTime(data.durationSec);
          setIsGenerating(false);
          break;

        case "error":
          console.error("Worker Error:", error);
          setErrorMsg(error);
          setModelStatus("error");
          setIsGenerating(false);
          break;

        default:
          break;
      }
    };

    workerRef.current.onmessage = onMessageReceived;

    // Auto-load model on mount (try WebGPU first)
    setModelStatus("loading");
    setProgressItems({});
    setTotalProgress(0);
    setErrorMsg("");
    setLoadTime(performance.now());
    workerRef.current.postMessage({ type: "load", data: { device: "wasm" } });

    // Clean up worker on unmount
    return () => {
      if (workerRef.current) {
        workerRef.current.terminate();
      }
    };
  }, []);

  function handleDeviceChange(newDevice) {
    if (isGenerating || modelStatus === "loading") return;
    setDeviceConfig(newDevice);
    setModelStatus("loading");
    setProgressItems({});
    setTotalProgress(0);
    setErrorMsg("");
    setLoadTime(performance.now());
    if (workerRef.current) {
      workerRef.current.postMessage({ type: "load", data: { device: newDevice } });
    }
  }

  async function handleClearCache() {
    if (window.confirm("Are you sure you want to delete the cached model files? This will require downloading the VLM weights again.")) {
      try {
        // Clear the custom IndexedDB cache
        await new Promise((resolve, reject) => {
          const req = indexedDB.deleteDatabase('transformers-cache');
          req.onsuccess = resolve;
          req.onerror = reject;
          req.onblocked = resolve;
        });

        // Clear fallback Cache API just in case
        const cacheKeys = await caches.keys();
        for (const key of cacheKeys) {
          if (key.includes("onnx") || key.includes("transformers") || key.includes("SmolVLM")) {
            await caches.delete(key);
          }
        }
        
        alert("Successfully cleared cached files from IndexedDB and Cache Storage!");
        window.location.reload();
      } catch (err) {
        alert("Failed to clear cache: " + err.message);
      }
    }
  }

  async function handleImageUpload(e) {
    const file = e.target.files[0];
    if (!file) return;

    setImageFile(file);
    setExtractedMetadata(null);
    setGenerationTime(null);

    if (file.type === "application/pdf") {
      try {
        const arrayBuffer = await file.arrayBuffer();
        const pdf = await pdfjsLib.getDocument({ data: arrayBuffer }).promise;
        const page = await pdf.getPage(1);
        const viewport = page.getViewport({ scale: 2.0 });
        
        const canvas = document.createElement("canvas");
        const context = canvas.getContext("2d");
        canvas.height = viewport.height;
        canvas.width = viewport.width;
        
        await page.render({ canvasContext: context, viewport }).promise;
        const dataUrl = canvas.toDataURL("image/png");
        setImagePreviewUrl(dataUrl);
      } catch (err) {
        console.error("PDF Parsing error:", err);
        alert("Failed to read PDF.");
      }
    } else {
      const reader = new FileReader();
      reader.onload = (ev) => {
        setImagePreviewUrl(ev.target.result);
      };
      reader.readAsDataURL(file);
    }
  }

  function handleExtractName() {
    if (!imagePreviewUrl || isGenerating || modelStatus !== "ready") return;

    setIsGenerating(true);
    setExtractedMetadata(null);
    setGenerationTime(null);

    const img = new Image();
    img.onload = () => {
      // Crop 1: Top Quarter (25%)
      const canvasQuarter = document.createElement("canvas");
      const ctxQuarter = canvasQuarter.getContext("2d");
      canvasQuarter.width = img.width;
      canvasQuarter.height = img.height * 0.25;
      ctxQuarter.drawImage(img, 0, 0, img.width, img.height * 0.25, 0, 0, canvasQuarter.width, canvasQuarter.height);
      // Crop: Top Half (50%) — enough context for both header info and subject deduction
      const canvasHalf = document.createElement("canvas");
      const ctxHalf = canvasHalf.getContext("2d");
      canvasHalf.width = img.width;
      canvasHalf.height = img.height * 0.50;
      ctxHalf.drawImage(img, 0, 0, img.width, img.height * 0.50, 0, 0, canvasHalf.width, canvasHalf.height);
      const croppedHalf = canvasHalf.toDataURL("image/png");

      if (workerRef.current) {
        workerRef.current.postMessage({
          type: "generate",
          data: {
            imageFileOrUrlTopHalf: croppedHalf,
            requestId: "vlm-extract"
          }
        });
      }
    };
    img.src = imagePreviewUrl;
  }

  function handleDownloadRenamed() {
    if (!imageFile || !extractedMetadata) return;

    // Clean name: remove all spaces, punctuation, leaving only alphanumeric
    const cleanName = (extractedMetadata.name || "Unknown")
      .replace(/[^a-zA-Z0-9]/g, "");
    
    // Clean roll number: remove any non-alphanumeric characters (like trailing periods)
    const cleanRoll = (extractedMetadata.rollNo || "")
      .replace(/[^a-zA-Z0-9]/g, "");

    // Map subjects to their clean abbreviations (Maths, Bio, Phy, Chem, Eco)
    const rawSubject = (extractedMetadata.subject || "").trim().toLowerCase();
    let subAbbr = "Sub";
    if (rawSubject.includes("math")) subAbbr = "Maths";
    else if (rawSubject.includes("bio")) subAbbr = "Bio";
    else if (rawSubject.includes("phys")) subAbbr = "Phy";
    else if (rawSubject.includes("chem")) subAbbr = "Chem";
    else if (rawSubject.includes("econ")) subAbbr = "Eco";
    else {
      // Fallback: clean the subject string and capitalize it
      const cleaned = rawSubject.replace(/[^a-z0-9]/gi, '');
      subAbbr = cleaned ? (cleaned.charAt(0).toUpperCase() + cleaned.slice(1)) : "Sub";
    }

    // Retain original extension (default to .pdf if missing)
    const extension = imageFile.name.includes('.') ? imageFile.name.split('.').pop() : 'pdf';
    
    // Format: StudentNameROLL_sub (e.g. AadarshPatel21_Maths.pdf)
    const newFilename = `${cleanName}${cleanRoll}_${subAbbr}.${extension}`;

    // Create download link
    const url = URL.createObjectURL(imageFile);
    const a = document.createElement('a');
    a.href = url;
    a.download = newFilename;
    document.body.appendChild(a);
    a.click();
    document.body.removeChild(a);
    URL.revokeObjectURL(url);
  }

  return (
    <div className="app-container">
      {/* Header Panel */}
      <div className="header-section">
        <div className="title-group">
          <h1>Student Name Extractor</h1>
          <p>Powered by SmolVLM-500M-Instruct running {modelStatus === "ready" && loadedDevice === "WebGPU" ? "on GPU via WebGPU" : "on CPU via WebAssembly"}</p>
          <div className="backend-control" style={{ marginTop: "10px", display: "flex", gap: "12px", alignItems: "center", flexWrap: "wrap" }}>
            <div style={{ display: "flex", gap: "6px", alignItems: "center" }}>
              <span style={{ fontSize: "13px", color: "rgba(255, 255, 255, 0.6)" }}>Device: <strong>WebAssembly (CPU)</strong></span>
            </div>

            <button
              onClick={handleClearCache}
              disabled={modelStatus === "loading" || isGenerating}
              style={{
                background: "rgba(239, 68, 68, 0.15)",
                border: "1px solid rgba(239, 68, 68, 0.3)",
                borderRadius: "6px",
                color: "#f87171",
                padding: "4px 8px",
                fontSize: "12px",
                cursor: "pointer",
                display: "flex",
                alignItems: "center",
                gap: "4px",
              }}
            >
              🗑️ Clear Storage Cache
            </button>
          </div>
        </div>
        <div className="model-badge">
          <span className={`pulse-dot ${modelStatus}`}></span>
          {modelStatus === "idle" && "Model Not Loaded"}
          {modelStatus === "loading" && `Downloading Weights: ${totalProgress.toFixed(0)}%`}
          {modelStatus === "ready" && `Model Ready (${loadedDevice})`}
          {modelStatus === "error" && "Load Failed"}
        </div>
      </div>


      {/* Main Content Area */}
      {modelStatus === "error" && (
        <div className="error-banner">
          ⚠️ <strong>Error loading model:</strong> {errorMsg}
          <br />
          <button onClick={() => window.location.reload()} style={{ marginTop: "10px", padding: "5px 10px", borderRadius: "4px", background: "rgba(255,255,255,0.1)", color: "white", border: "none", cursor: "pointer" }}>Retry Loading</button>
        </div>
      )}

      {modelStatus !== "ready" && modelStatus !== "error" && (
        <div className="glass-panel" style={{ textAlign: "center", padding: "60px 20px" }}>
          <h2>Initializing VLM Scanner Engine</h2>
          <p style={{ color: "rgba(255,255,255,0.7)", marginTop: "10px", marginBottom: "30px" }}>
            Downloading and compiling the SmolVLM-500M-Instruct model. <br />
            This will take a few minutes on the first run as it caches the ~800MB weights.
          </p>
          <div className="progress-container" style={{ width: "80%", margin: "0 auto", height: "12px", background: "rgba(255,255,255,0.1)", borderRadius: "6px", overflow: "hidden" }}>
            <div className="progress-bar" style={{ width: `${totalProgress}%`, height: "100%", background: "var(--primary)", transition: "width 0.3s ease" }}></div>
          </div>
        </div>
      )}

      {modelStatus === "ready" && (
        <div className="vlm-scanner-container" style={{ display: "grid", gridTemplateColumns: "1fr 1fr", gap: "20px" }}>
          
          {/* Upload Panel */}
          <div className="glass-panel" style={{ display: "flex", flexDirection: "column", gap: "20px" }}>
            <h3>1. Upload Document</h3>
            <p style={{ fontSize: "14px", color: "rgba(255,255,255,0.6)" }}>
              Upload an image of a document (PNG, JPG, WebP) containing a student's name.
            </p>
            
            <div 
              style={{
                border: "2px dashed rgba(255,255,255,0.2)",
                borderRadius: "12px",
                padding: "40px",
                textAlign: "center",
                position: "relative",
                cursor: "pointer",
                transition: "all 0.2s ease",
                backgroundColor: "rgba(0,0,0,0.2)"
              }}
              onClick={() => document.getElementById("file-upload").click()}
            >
              <input 
                id="file-upload" 
                type="file" 
                accept="image/*,application/pdf" 
                onChange={handleImageUpload} 
                style={{ display: "none" }} 
              />
              <span style={{ fontSize: "32px", display: "block", marginBottom: "10px" }}>📄</span>
              <span style={{ fontWeight: "500" }}>Click to select a document</span>
            </div>

              {imagePreviewUrl && (
              <button 
                onClick={handleExtractName} 
                disabled={isGenerating}
                className="send-button"
                style={{ padding: "14px", width: "100%", fontSize: "16px", borderRadius: "8px", marginTop: "auto" }}
              >
                {isGenerating ? "Scanning Document..." : "Extract Metadata"}
              </button>
            )}
          </div>

          {/* Preview & Results Panel */}
          <div style={{ display: "flex", flexDirection: "column", gap: "20px" }}>
            
            <div className="glass-panel" style={{ flex: 1, display: "flex", flexDirection: "column" }}>
              <div style={{ display: "flex", justifyContent: "space-between", alignItems: "center", marginBottom: "15px" }}>
                <h3 style={{ margin: 0 }}>Document Preview</h3>
              </div>
              
              <div style={{ 
                flex: 1, 
                backgroundColor: "rgba(0,0,0,0.3)", 
                borderRadius: "8px", 
                display: "flex", 
                alignItems: "center", 
                justifyContent: "center",
                overflow: "hidden",
                minHeight: "250px"
              }}>
                {imagePreviewUrl ? (
                  <img src={imagePreviewUrl} alt="Document Preview" style={{ maxWidth: "100%", maxHeight: "300px", objectFit: "contain" }} />
                ) : (
                  <span style={{ color: "rgba(255,255,255,0.3)", fontSize: "14px" }}>No document selected</span>
                )}
              </div>
            </div>

            {/* Document Metadata Panel */}
            <div className="glass-panel" style={{ flex: 1, display: "flex", flexDirection: "column" }}>
              <div style={{ display: "flex", justifyContent: "space-between", alignItems: "center", marginBottom: "15px" }}>
                <h3 style={{ margin: 0, fontSize: "14px", letterSpacing: "1px", color: "rgba(255,255,255,0.6)" }}>EXTRACTED METADATA</h3>
                <div style={{ display: "flex", alignItems: "center", gap: "8px" }}>
                  {generationTime && (
                    <span style={{ fontSize: "12px", color: "rgba(255, 255, 255, 0.4)" }}>
                      Processed in {generationTime}s
                    </span>
                  )}
                  {isGenerating && <span className="pulse-dot"></span>}
                </div>
              </div>
              
              <div style={{ 
                padding: "20px", 
                backgroundColor: "rgba(0,0,0,0.2)", 
                borderRadius: "8px",
                border: "1px solid rgba(255,255,255,0.05)",
                textAlign: "left",
                flex: 1,
                display: "flex",
                flexDirection: "column",
                gap: "15px"
              }}>
                {isGenerating ? (
                   <div style={{ flex: 1, display: "flex", alignItems: "center", justifyContent: "center", color: "rgba(255,255,255,0.4)" }}>
                     Analyzing document structure...
                   </div>
                ) : extractedMetadata ? (
                   <>
                     <div style={{ display: "flex", flexDirection: "column", gap: "5px" }}>
                        <span style={{ fontSize: "12px", color: "rgba(255,255,255,0.5)", textTransform: "uppercase" }}>Student Name</span>
                        <strong style={{ fontSize: "18px", color: "var(--primary)" }}>{extractedMetadata.name || "--"}</strong>
                     </div>
                     <div style={{ display: "flex", flexDirection: "column", gap: "5px" }}>
                        <span style={{ fontSize: "12px", color: "rgba(255,255,255,0.5)", textTransform: "uppercase" }}>Subject / Course</span>
                        <strong style={{ fontSize: "18px", color: "white" }}>{extractedMetadata.subject || "--"}</strong>
                     </div>
                     <div style={{ display: "flex", flexDirection: "column", gap: "5px" }}>
                        <span style={{ fontSize: "12px", color: "rgba(255,255,255,0.5)", textTransform: "uppercase" }}>Roll Number / ID</span>
                        <strong style={{ fontSize: "18px", color: "white" }}>{extractedMetadata.rollNo || "--"}</strong>
                     </div>
                     
                     <div style={{ marginTop: "auto", paddingTop: "20px" }}>
                       <button 
                         onClick={handleDownloadRenamed} 
                         style={{ 
                           width: "100%", 
                           padding: "12px", 
                           backgroundColor: "rgba(16, 185, 129, 0.2)", 
                           color: "#34d399", 
                           border: "1px solid rgba(16, 185, 129, 0.4)", 
                           borderRadius: "6px", 
                           cursor: "pointer",
                           fontWeight: "bold",
                           display: "flex",
                           justifyContent: "center",
                           alignItems: "center",
                           gap: "8px"
                         }}
                       >
                         ⬇️ Download Renamed File
                       </button>
                     </div>
                   </>
                ) : (
                   <div style={{ flex: 1, display: "flex", alignItems: "center", justifyContent: "center", color: "rgba(255,255,255,0.4)" }}>
                     Metadata will appear here
                   </div>
                )}
              </div>
            </div>

          </div>
        </div>
      )}
    </div>
  );
}

export default App;