import React, { useState, useEffect, useRef } from 'react';
import { initializeApp } from "firebase/app";
import { getAuth, signInAnonymously, onAuthStateChanged, signInWithCustomToken } from "firebase/auth";
import { getFirestore, collection, addDoc, getDocs, query, where } from "firebase/firestore";
import { 
  Camera, Upload, User, CheckCircle, XCircle, Loader2, ScanFace, 
  LogIn, LayoutDashboard, LogOut, ChevronLeft, 
  History, Download, ShieldCheck, RefreshCw, AlertCircle
} from 'lucide-react';

// --- Constants ---
const MODEL_URL = 'https://justadudewhohacks.github.io/face-api.js/models';
const ADMIN_PASSCODE = "admin123"; 
const MATCH_THRESHOLD = 0.5; // Strictness (Lower is stricter)

const FaceAttendanceApp = () => {
  // --- Global State ---
  const [firebaseInit, setFirebaseInit] = useState(false);
  const [modelsLoaded, setModelsLoaded] = useState(false);
  
  // --- Data State ---
  const [studentProfile, setStudentProfile] = useState(null); 
  const [attendanceHistory, setAttendanceHistory] = useState([]);
  const [adminStudents, setAdminStudents] = useState([]);
  
  // --- UI State ---
  const [view, setView] = useState('landing'); 
  const [status, setStatus] = useState({ type: '', msg: '' });
  const [isProcessing, setIsProcessing] = useState(false);
  const [selectedStudent, setSelectedStudent] = useState(null);

  // --- Input Method ---
  const [method, setMethod] = useState('camera'); 
  const [imageSrc, setImageSrc] = useState(null); 

  // --- Refs ---
  const videoRef = useRef(null);
  const imgRef = useRef(null);
  const streamRef = useRef(null);
  const canvasRef = useRef(null); // Hidden canvas for processing

  // --- Firebase ---
  const [db, setDb] = useState(null);
  const [appId, setAppId] = useState(null);

  // 1. Initialize Firebase
  useEffect(() => {
    try {
        const config = JSON.parse(__firebase_config);
        const app = initializeApp(config);
        const auth = getAuth(app);
        const firestore = getFirestore(app);
        setDb(firestore);
        setAppId(typeof __app_id !== 'undefined' ? __app_id : 'default-app-id');

        const initAuth = async () => {
            if (typeof __initial_auth_token !== 'undefined' && __initial_auth_token) {
                await signInWithCustomToken(auth, __initial_auth_token);
            } else {
                await signInAnonymously(auth);
            }
        };
        initAuth();
        onAuthStateChanged(auth, () => setFirebaseInit(true));
    } catch (e) {
        console.error("Firebase init error:", e);
    }
  }, []);

  // 2. Load Scripts & Models
  useEffect(() => {
    const loadScripts = async () => {
      if (!window.faceapi) {
        const script = document.createElement('script');
        script.src = 'https://cdn.jsdelivr.net/npm/face-api.js@0.22.2/dist/face-api.min.js';
        script.async = true;
        script.onload = () => loadModels();
        document.body.appendChild(script);
      } else {
        loadModels();
      }

      if (!window.jspdf) {
         const pdfScript = document.createElement('script');
         pdfScript.src = 'https://cdnjs.cloudflare.com/ajax/libs/jspdf/2.5.1/jspdf.umd.min.js';
         document.body.appendChild(pdfScript);
         const tableScript = document.createElement('script');
         tableScript.src = 'https://cdnjs.cloudflare.com/ajax/libs/jspdf-autotable/3.5.29/jspdf.plugin.autotable.min.js';
         document.body.appendChild(tableScript);
      }
    };

    const loadModels = async () => {
      try {
        // Load TinyFaceDetector (Much lighter and faster than SSD)
        await Promise.all([
          window.faceapi.nets.tinyFaceDetector.loadFromUri(MODEL_URL),
          window.faceapi.nets.faceLandmark68Net.loadFromUri(MODEL_URL),
          window.faceapi.nets.faceRecognitionNet.loadFromUri(MODEL_URL)
        ]);
        setModelsLoaded(true);
      } catch (err) {
        setStatus({ type: 'error', msg: 'AI Models failed. Reload page.' });
      }
    };
    loadScripts();
  }, []);

  // 3. Clear Status Timer
  useEffect(() => {
    if (status.msg && status.type !== 'loading') {
        const timer = setTimeout(() => setStatus({ type: '', msg: '' }), 4000);
        return () => clearTimeout(timer);
    }
  }, [status]);

  useEffect(() => {
    setStatus({ type: '', msg: '' });
    setIsProcessing(false);
    setImageSrc(null);
    setMethod('camera');
  }, [view]);

  // --- Camera Logic ---

  const startVideo = async () => {
    if (!modelsLoaded || method !== 'camera') return;
    try {
      const stream = await navigator.mediaDevices.getUserMedia({ 
          video: { facingMode: 'user', width: { ideal: 640 }, height: { ideal: 480 } } 
      });
      if (videoRef.current) {
        videoRef.current.srcObject = stream;
        streamRef.current = stream;
        // Important: Wait for metadata to load to prevent size errors
        videoRef.current.onloadedmetadata = () => {
             videoRef.current.play().catch(e => console.log("Play error", e));
        };
      }
    } catch (err) {
      setStatus({ type: 'error', msg: 'Camera blocked. Check permissions.' });
    }
  };

  const stopVideo = () => {
    if (streamRef.current) {
      streamRef.current.getTracks().forEach(track => track.stop());
      streamRef.current = null;
    }
    if (videoRef.current) videoRef.current.srcObject = null;
  };

  useEffect(() => {
    if (['register', 'login', 'mark-attendance'].includes(view) && method === 'camera') {
      startVideo();
    } else {
      stopVideo();
    }
    return () => stopVideo();
  }, [view, method, modelsLoaded]);

  const handleFileChange = (e) => {
      if (e.target.files && e.target.files[0]) {
          const file = e.target.files[0];
          const reader = new FileReader();
          reader.onload = (ev) => setImageSrc(ev.target.result);
          reader.readAsDataURL(file);
      }
  };

  // --- STABLE AI DETECTION (The Fix) ---

  const extractFaceFromVideo = async () => {
      const video = videoRef.current;
      if (!video || video.readyState !== 4) return null;

      // 1. Draw video frame to hidden canvas first
      // This isolates the AI from the video stream, preventing freezes
      const canvas = document.createElement('canvas');
      canvas.width = video.videoWidth;
      canvas.height = video.videoHeight;
      const ctx = canvas.getContext('2d');
      ctx.drawImage(video, 0, 0, canvas.width, canvas.height);

      // 2. Run detection on the STATIC canvas, not the live video
      // Use TinyFaceDetectorOptions for speed
      const detection = await window.faceapi.detectSingleFace(
          canvas, 
          new window.faceapi.TinyFaceDetectorOptions()
      ).withFaceLandmarks().withFaceDescriptor();

      return detection;
  };

  const extractFaceFromImage = async () => {
      const img = imgRef.current;
      if (!img) return null;
      
      const detection = await window.faceapi.detectSingleFace(
          img,
          new window.faceapi.TinyFaceDetectorOptions()
      ).withFaceLandmarks().withFaceDescriptor();

      return detection;
  };

  const getReliableDetection = async () => {
      // Timeout protection: If detection takes > 5s, force stop
      const timeoutPromise = new Promise((_, reject) => 
          setTimeout(() => reject(new Error("Detection timed out. Move closer/better light.")), 5000)
      );

      const detectionLogic = async () => {
          if (method === 'camera') return await extractFaceFromVideo();
          return await extractFaceFromImage();
      };

      try {
          return await Promise.race([detectionLogic(), timeoutPromise]);
      } catch (error) {
          throw error;
      }
  };

  const captureSnapshot = () => {
      const canvas = document.createElement('canvas');
      let input, width, height;

      if (method === 'camera' && videoRef.current) {
          input = videoRef.current;
          width = input.videoWidth;
          height = input.videoHeight;
      } else if (method === 'upload' && imgRef.current) {
          input = imgRef.current;
          width = input.naturalWidth;
          height = input.naturalHeight;
      }

      if (!input || !width || !height) return null;

      const maxDim = 600;
      const scale = Math.min(maxDim / width, maxDim / height, 1);
      canvas.width = width * scale;
      canvas.height = height * scale;
      
      const ctx = canvas.getContext('2d');
      if (method === 'camera') {
          ctx.translate(canvas.width, 0);
          ctx.scale(-1, 1);
      }
      ctx.drawImage(input, 0, 0, canvas.width, canvas.height);
      return canvas.toDataURL('image/jpeg', 0.85);
  };

  // --- HANDLERS ---

  const handleRegister = async (name) => {
    if (!name.trim()) return setStatus({type: 'error', msg: 'Name required'});
    if (method === 'upload' && !imageSrc) return setStatus({type: 'error', msg: 'Upload an image'});

    setIsProcessing(true);
    setStatus({ type: 'loading', msg: 'Processing Face...' });

    try {
      const detection = await getReliableDetection();
      if (!detection) throw new Error("No face found. Hold still & ensure good light.");

      const descriptor = Array.from(detection.descriptor); 
      const photoUrl = captureSnapshot(); 

      const docRef = await addDoc(collection(db, 'artifacts', appId, 'public', 'data', 'students'), {
        name: name,
        descriptor: descriptor,
        photoUrl: photoUrl,
        registeredAt: new Date().toISOString()
      });

      setStudentProfile({ id: docRef.id, name, photoUrl, descriptor });
      setView('student-dashboard');
      setStatus({ type: 'success', msg: 'Registration Complete!' });

    } catch (err) {
      setStatus({ type: 'error', msg: err.message || "Failed. Try again." });
    } finally {
        setIsProcessing(false);
    }
  };

  const handleLogin = async () => {
    if (method === 'upload' && !imageSrc) return setStatus({type: 'error', msg: 'Upload an image'});
    
    setIsProcessing(true);
    setStatus({ type: 'loading', msg: 'Scanning Database...' });

    try {
        const q = query(collection(db, 'artifacts', appId, 'public', 'data', 'students'));
        const querySnapshot = await getDocs(q);
        
        const labeledDescriptors = [];
        const studentsMap = {};

        querySnapshot.forEach((doc) => {
            const data = doc.data();
            if (data.descriptor && data.descriptor.length > 0) {
                const descArray = new Float32Array(data.descriptor);
                labeledDescriptors.push(new window.faceapi.LabeledFaceDescriptors(doc.id, [descArray]));
                studentsMap[doc.id] = { id: doc.id, ...data };
            }
        });

        if (labeledDescriptors.length === 0) throw new Error("No students registered.");

        setStatus({ type: 'loading', msg: 'Analyzing Face...' });
        const detection = await getReliableDetection();

        if (!detection) throw new Error("No face detected.");

        const faceMatcher = new window.faceapi.FaceMatcher(labeledDescriptors, MATCH_THRESHOLD);
        const bestMatch = faceMatcher.findBestMatch(detection.descriptor);

        if (bestMatch.label === 'unknown') throw new Error("Face not recognized. Please Register.");

        const student = studentsMap[bestMatch.label];
        setStudentProfile(student);
        await fetchAttendanceHistory(student.id);
        setView('student-dashboard');
        setStatus({ type: 'success', msg: `Welcome back, ${student.name}!` });

    } catch (err) {
        setStatus({ type: 'error', msg: err.message });
    } finally {
        setIsProcessing(false);
    }
  };

  const handleMarkAttendance = async () => {
      if (method === 'upload' && !imageSrc) return setStatus({type: 'error', msg: 'Upload an image'});
      
      setIsProcessing(true);
      setStatus({ type: 'loading', msg: 'Verifying You...' });

      try {
        if (!studentProfile) throw new Error("Session expired. Login again.");

        const detection = await getReliableDetection();
        
        if (!detection) throw new Error("No face detected.");

        // Strict Check
        const currentUserDescriptor = new window.faceapi.LabeledFaceDescriptors(
            studentProfile.id,
            [new Float32Array(studentProfile.descriptor)]
        );
        
        const faceMatcher = new window.faceapi.FaceMatcher([currentUserDescriptor], MATCH_THRESHOLD);
        const match = faceMatcher.findBestMatch(detection.descriptor);

        if (match.label === 'unknown') {
            throw new Error("Face Mismatch! Not the logged-in student.");
        }

        const proofPhoto = captureSnapshot();
        const now = new Date();
        
        await addDoc(collection(db, 'artifacts', appId, 'public', 'data', 'attendance'), {
            studentId: studentProfile.id,
            studentName: studentProfile.name,
            timestamp: now.toISOString(),
            date: now.toLocaleDateString(), 
            proof: proofPhoto
        });

        setStatus({ type: 'success', msg: 'Marked Present!' });
        await fetchAttendanceHistory(studentProfile.id);
        
        setTimeout(() => setView('student-dashboard'), 1500);

      } catch (err) {
          setStatus({ type: 'error', msg: err.message });
      } finally {
          setIsProcessing(false);
      }
  };

  // --- HELPERS ---

  const fetchAttendanceHistory = async (studentId) => {
      const q = query(
          collection(db, 'artifacts', appId, 'public', 'data', 'attendance'),
          where('studentId', '==', studentId)
      );
      const snap = await getDocs(q);
      const records = snap.docs.map(d => ({id: d.id, ...d.data()}));
      records.sort((a,b) => new Date(b.timestamp) - new Date(a.timestamp));
      setAttendanceHistory(records);
  };

  const fetchAllStudents = async () => {
      setIsProcessing(true);
      const q = query(collection(db, 'artifacts', appId, 'public', 'data', 'students'));
      const snap = await getDocs(q);
      setAdminStudents(snap.docs.map(d => ({id: d.id, ...d.data()})));
      setIsProcessing(false);
  };

  const isAttendanceMarkedToday = () => {
      const today = new Date().toLocaleDateString();
      return attendanceHistory.some(rec => rec.date === today);
  };

  const generatePDF = (student, history) => {
      if (!window.jspdf) return setStatus({type: 'error', msg: 'PDF Tool loading...'});
      const { jsPDF } = window.jspdf;
      const doc = new jsPDF();
      doc.text(`Attendance Report: ${student.name}`, 14, 20);
      const rows = history.map(h => [new Date(h.timestamp).toLocaleDateString(), new Date(h.timestamp).toLocaleTimeString(), "Present"]);
      doc.autoTable({ startY: 30, head: [['Date', 'Time', 'Status']], body: rows });
      doc.save(`${student.name}_Report.pdf`);
  };

  // --- RENDER ---

  const LoadingOverlay = () => (
      <div className="fixed inset-0 bg-black/80 flex items-center justify-center z-50 backdrop-blur-sm">
          <div className="bg-slate-900 p-6 rounded-2xl flex flex-col items-center border border-slate-700">
              <Loader2 className="w-10 h-10 text-blue-500 animate-spin mb-3"/>
              <p className="text-white font-medium">{status.msg || 'Processing...'}</p>
          </div>
      </div>
  );

  const InputSwitcher = () => (
      <div className="flex bg-slate-900 p-1 rounded-lg border border-slate-800 mb-4 w-full max-w-xs mx-auto">
          <button 
              onClick={() => { setMethod('camera'); setImageSrc(null); }}
              className={`flex-1 py-2 text-sm font-medium rounded-md flex items-center justify-center gap-2 transition-all ${method === 'camera' ? 'bg-blue-600 text-white shadow' : 'text-slate-400 hover:text-white'}`}
          >
              <Camera size={16} /> Live Camera
          </button>
          <button 
              onClick={() => { setMethod('upload'); setImageSrc(null); }}
              className={`flex-1 py-2 text-sm font-medium rounded-md flex items-center justify-center gap-2 transition-all ${method === 'upload' ? 'bg-indigo-600 text-white shadow' : 'text-slate-400 hover:text-white'}`}
          >
              <Upload size={16} /> Upload Photo
          </button>
      </div>
  );

  const MediaDisplay = ({ title, btnText, action }) => (
    <div className="flex flex-col items-center w-full max-w-md mx-auto animate-in fade-in">
        <h2 className="text-xl font-bold text-white mb-4">{title}</h2>
        <InputSwitcher />
        <div className="relative w-full aspect-video bg-black rounded-2xl overflow-hidden border border-slate-700 shadow-2xl mb-6 group flex items-center justify-center">
            {method === 'camera' ? (
                <>
                    <video ref={videoRef} autoPlay muted playsInline className="w-full h-full object-cover transform scale-x-[-1]" />
                    <div className="absolute top-0 left-0 w-full h-1 bg-blue-500/50 shadow-[0_0_15px_rgba(59,130,246,0.5)] animate-[scan_2s_ease-in-out_infinite]"></div>
                </>
            ) : (
                <>
                    {imageSrc ? (
                        <img ref={imgRef} src={imageSrc} className="w-full h-full object-contain" alt="preview" />
                    ) : (
                        <label className="flex flex-col items-center justify-center w-full h-full cursor-pointer hover:bg-slate-900 transition-colors">
                            <Upload size={40} className="text-slate-600 mb-2" />
                            <span className="text-slate-500 text-sm">Click to upload image</span>
                            <input type="file" accept="image/*" onChange={handleFileChange} className="hidden" />
                        </label>
                    )}
                </>
            )}
        </div>
        <button 
            onClick={action}
            disabled={isProcessing || (method === 'upload' && !imageSrc)}
            className="w-full py-4 bg-gradient-to-r from-blue-600 to-indigo-600 hover:from-blue-500 hover:to-indigo-500 text-white rounded-xl font-bold shadow-lg shadow-blue-500/20 transition-all active:scale-95 disabled:opacity-50 disabled:cursor-not-allowed"
        >
            {btnText}
        </button>
    </div>
  );

  // --- APP STRUCTURE ---

  if (!modelsLoaded || !firebaseInit) {
      return (
          <div className="min-h-screen bg-slate-950 flex flex-col items-center justify-center text-white">
              <Loader2 className="w-12 h-12 text-blue-500 animate-spin mb-4" />
              <p className="text-slate-400">Initializing Core Systems...</p>
          </div>
      );
  }

  return (
    <div className="min-h-screen bg-slate-950 text-slate-200 font-sans pb-10">
      <header className="p-5 border-b border-slate-800 bg-slate-900/50 backdrop-blur sticky top-0 z-40">
          <div className="max-w-5xl mx-auto flex items-center justify-between">
              <div className="flex items-center gap-3 cursor-pointer" onClick={() => setView('landing')}>
                  <div className="w-10 h-10 bg-gradient-to-br from-blue-600 to-indigo-600 rounded-lg flex items-center justify-center shadow-lg">
                      <ScanFace className="text-white w-6 h-6" />
                  </div>
                  <h1 className="font-bold text-lg text-white">FaceAuth Pro</h1>
              </div>
              {view !== 'landing' && <button onClick={() => setView('landing')} className="text-sm text-slate-400 hover:text-white">Home</button>}
          </div>
      </header>

      <main className="max-w-5xl mx-auto p-4 mt-6">
          {view === 'landing' && (
              <div className="flex flex-col items-center justify-center min-h-[60vh] text-center space-y-8 animate-in zoom-in-95">
                  <h1 className="text-5xl font-extrabold text-white">Smart Attendance</h1>
                  <div className="grid grid-cols-1 sm:grid-cols-2 gap-4 w-full max-w-md">
                      <button onClick={() =