import React, { useState, Suspense } from 'react';
import { useNavigate } from 'react-router-dom';
import { Sparkles, Upload, FileText, Image as ImageIcon, Box, Check, Loader2, ArrowRight, RefreshCw, ChevronLeft } from 'lucide-react';
import { api } from '../lib/api';
import { Canvas } from '@react-three/fiber';
import { OrbitControls, Stage, useGLTF } from '@react-three/drei';

function ModelPreview({ url }: { url: string }) {
  const { scene } = useGLTF(url);
  return <primitive object={scene} />;
}

const STEPS = [
  { id: 1, title: 'Upload PDF', icon: Upload },
  { id: 2, title: 'AI Review', icon: FileText },
  { id: 3, title: 'Concept Art', icon: ImageIcon },
  { id: 4, title: '3D Generation', icon: Box },
];

export default function SparkPlug() {
  const navigate = useNavigate();
  const [currentStep, setCurrentStep] = useState(1);
  const [isLoading, setIsLoading] = useState(false);
  const [loadingMessage, setLoadingMessage] = useState('');
  
  // Step 1: PDF File
  const [pdfFile, setPdfFile] = useState<File | null>(null);

  // Step 2: Extracted Text
  const [extractedPrompt, setExtractedPrompt] = useState<string>('');

  // Step 3: Generated Images (Multi-view)
  const [generatedImages, setGeneratedImages] = useState<string[]>([]);

  // Step 4: 3D Model
  const [modelProgress, setModelProgress] = useState(0);
  const [generatedGlbUrl, setGeneratedGlbUrl] = useState<string | null>(null);

  const handleFileUpload = async (e: React.ChangeEvent<HTMLInputElement>) => {
    if (e.target.files && e.target.files[0]) {
      setPdfFile(e.target.files[0]);
    }
  };

  const processStep1 = async () => {
    if (!pdfFile) return;
    setLoadingMessage('Extracting details with Gemini AI...');
    setIsLoading(true);
    
    try {
      // Validate keys
      const keys = await api.getUserKeys();
      const geminiKey = keys.GEMINI_API_KEY;
      
      if (!geminiKey) {
        alert("Please configure your Gemini API Key in the Dashboard Settings first.");
        setIsLoading(false);
        return;
      }
      // Call our backend endpoint to parse PDF and query Gemini
      const prompt = await api.extractPDF(pdfFile);
      setExtractedPrompt(prompt);
      
      setCurrentStep(2);
    } catch (err: any) {
      console.error(err);
      alert(`Failed to process PDF: ${err.message || 'Unknown error'}`);
    } finally {
      setIsLoading(false);
    }
  };

  const processStep2 = async () => {
    if (!extractedPrompt) return;
    setLoadingMessage('Generating multi-view concepts with Nano Banana...');
    setIsLoading(true);
    
    try {
      // Validate keys
      const keys = await api.getUserKeys();
      const nanoBananaKey = keys.NANO_BANANA_API_KEY;
      
      if (!nanoBananaKey) {
        alert("Please configure your Nano Banana API Key in the Dashboard Settings first.");
        setIsLoading(false);
        return;
      }
      // Call our backend endpoint which queries Gemini Imagen
      const images = await api.generateImages(extractedPrompt);
      if (images && images.length > 0) {
        setGeneratedImages(images);
        setCurrentStep(3);
      } else {
        throw new Error("Failed to generate images.");
      }
    } catch (err: any) {
      console.error(err);
      alert(`Failed to generate images: ${err.message || 'Unknown error'}`);
    } finally {
      setIsLoading(false);
    }
  };

  const processStep3 = async () => {
    if (generatedImages.length === 0) return;
    setLoadingMessage('Constructing 3D Model with Meshy API...');
    setIsLoading(true);
    
    try {
      // Validate keys
      const keys = await api.getUserKeys();
      const meshyKey = keys.MESHY_API_KEY;
      
      if (!meshyKey) {
        alert("Please configure your Meshy API Key in the Dashboard Settings first.");
        setIsLoading(false);
        return;
      }

      // Call backend to start Meshy task
      const taskId = await api.generate3D(generatedImages);
      
      // Immediately move to step 4 to show the progress bar
      setCurrentStep(4);
      setIsLoading(false);
      
      // Poll Meshy status
      let isDone = false;
      let attempts = 0;
      while (!isDone && attempts < 60) { // Max 5 mins (60 * 5s)
        attempts++;
        await new Promise(resolve => setTimeout(resolve, 5000));
        
        const statusData = await api.get3DStatus(taskId);
        const { status, progress, model_urls, task_error } = statusData;
        
        if (progress) {
          setModelProgress(progress);
        }
        
        if (status === 'SUCCEEDED') {
          setModelProgress(100);
          setGeneratedGlbUrl(model_urls.glb);
          isDone = true;
          // The step is already 4, user just clicks "Open in Editor" when progress >= 100
        } else if (status === 'FAILED') {
          throw new Error(`Meshy generation failed: ${task_error?.message || 'Unknown error'}`);
        }
      }
      
      if (!isDone) {
        throw new Error("Meshy task timed out after 5 minutes.");
      }
      
    } catch (err: any) {
      console.error(err);
      alert(`Failed to construct 3D model: ${err.message || 'Unknown error'}`);
      // If it fails, revert back to step 3 so they can try again
      setCurrentStep(3);
    } finally {
      setIsLoading(false);
    }
  };

  const processStep4 = async () => {
    setLoadingMessage('Saving project to your workspace...');
    setIsLoading(true);
    try {
      if (!generatedGlbUrl) {
        throw new Error("No GLB URL available to save.");
      }
      
      // Create project and save with model injected into scene_data
      const res = await api.saveProject({
        name: pdfFile?.name ? pdfFile.name.replace(/\.[^/.]+$/, "") : "Spark Plug Output",
        scene_data: {
          sceneObjects: [],
          modelData: {
            name: "Generated Product Model",
            data: generatedGlbUrl,
            extension: "glb"
          },
          backgroundColor: '#444444',
          backgroundType: 'color',
          envPreset: 'studio',
          lightProps: { intensity: 1, color: '#ffffff', shadowBias: -0.0005, shadowRadius: 4 },
          groundProps: { showShadow: true, shadowIntensity: 0.4, shadowRotation: 0, shadowSoftness: 0.5, shadowLength: 10 }
        }
      });
      if (res.id) {
        navigate(`/editor/${res.id}`);
      }
    } catch (err) {
      console.error('Failed to create project', err);
      alert('Failed to save project');
    } finally {
      setIsLoading(false);
    }
  };

  return (
    <div className="min-h-screen bg-zinc-950 text-white flex flex-col">
      {/* Header */}
      <header className="h-16 border-b border-zinc-800 bg-zinc-900 flex items-center px-8 shrink-0">
        <button 
          onClick={() => navigate('/dashboard')}
          className="flex items-center gap-2 text-zinc-400 hover:text-white transition-colors mr-6"
        >
          <ChevronLeft size={18} />
          <span className="text-sm font-medium">Dashboard</span>
        </button>
        <div className="flex items-center gap-2 font-bold text-lg tracking-tight">
          <Sparkles className="text-amber-500" size={20} />
          Spark Plug <span className="text-zinc-500 font-normal text-sm ml-2">| Idea to 3D Generation</span>
        </div>
      </header>

      <main className="flex-1 flex flex-col items-center max-w-5xl mx-auto w-full p-8">
        
        {/* Stepper */}
        <div className="w-full flex items-center justify-between mb-12 relative">
          <div className="absolute left-0 top-1/2 -translate-y-1/2 w-full h-0.5 bg-zinc-800 -z-10" />
          {STEPS.map((step) => {
            const Icon = step.icon;
            const isCompleted = currentStep > step.id;
            const isCurrent = currentStep === step.id;
            
            return (
              <div key={step.id} className="flex flex-col items-center gap-2 bg-zinc-950 px-4">
                <div className={`w-10 h-10 rounded-full flex items-center justify-center border-2 transition-colors ${
                  isCompleted ? 'bg-amber-500 border-amber-500 text-white' : 
                  isCurrent ? 'bg-zinc-900 border-amber-500 text-amber-500' : 
                  'bg-zinc-900 border-zinc-800 text-zinc-600'
                }`}>
                  {isCompleted ? <Check size={18} /> : <Icon size={18} />}
                </div>
                <span className={`text-xs font-medium ${isCurrent ? 'text-amber-500' : isCompleted ? 'text-zinc-300' : 'text-zinc-600'}`}>
                  {step.title}
                </span>
              </div>
            );
          })}
        </div>

        {/* Content Area */}
        <div className="w-full max-w-3xl bg-zinc-900 border border-zinc-800 rounded-2xl p-8 shadow-2xl min-h-[400px] flex flex-col relative overflow-hidden">
          
          {isLoading && (
            <div className="absolute inset-0 z-50 bg-zinc-900/80 backdrop-blur-sm flex flex-col items-center justify-center animate-in fade-in duration-300">
              <div className="w-16 h-16 rounded-full border-4 border-zinc-800 border-t-amber-500 animate-spin mb-6 flex items-center justify-center">
                <Sparkles className="text-amber-500 animate-pulse" size={24} />
              </div>
              <h3 className="text-lg font-bold text-white mb-2">{loadingMessage}</h3>
              <p className="text-sm text-zinc-400">This might take a few moments...</p>
            </div>
          )}

          {currentStep === 1 && (
            <div className="flex-1 flex flex-col items-center justify-center gap-6 animate-in fade-in zoom-in duration-300">
              <div className="w-20 h-20 bg-zinc-800 rounded-full flex items-center justify-center">
                <Upload size={32} className="text-zinc-400" />
              </div>
              <div className="text-center">
                <h2 className="text-xl font-bold mb-2">Upload Mintel Spark PDF</h2>
                <p className="text-sm text-zinc-400 max-w-md mx-auto">
                  Upload a product concept PDF. Spark Plug will extract the product details using Gemini AI to prepare for 3D generation.
                </p>
              </div>
              
              <div className="relative group w-full max-w-md mt-4">
                <input 
                  type="file" 
                  accept=".pdf"
                  onChange={handleFileUpload}
                  className="absolute inset-0 w-full h-full opacity-0 cursor-pointer z-10"
                />
                <div className={`w-full p-6 border-2 border-dashed rounded-xl flex flex-col items-center gap-3 transition-colors ${pdfFile ? 'border-amber-500 bg-amber-500/10' : 'border-zinc-700 bg-zinc-950 group-hover:border-zinc-500'}`}>
                  <FileText size={24} className={pdfFile ? 'text-amber-500' : 'text-zinc-500'} />
                  <span className="text-sm font-medium text-zinc-300">
                    {pdfFile ? pdfFile.name : 'Click or drag PDF here'}
                  </span>
                </div>
              </div>

              {pdfFile && (
                <button 
                  onClick={processStep1}
                  disabled={isLoading}
                  className="mt-4 flex items-center gap-2 bg-amber-500 hover:bg-amber-600 text-black px-6 py-3 rounded-xl font-bold transition-all disabled:opacity-50"
                >
                  {isLoading ? <Loader2 className="animate-spin" size={18} /> : <Sparkles size={18} />}
                  Extract with Gemini
                </button>
              )}
            </div>
          )}

          {currentStep === 2 && (
            <div className="flex-1 flex flex-col gap-4 animate-in fade-in slide-in-from-right-8 duration-300">
              <h2 className="text-xl font-bold">Review Extracted Concept</h2>
              <p className="text-sm text-zinc-400">
                Gemini AI has analyzed the PDF. Review and modify the prompt below before generating images.
              </p>

              <textarea 
                value={extractedPrompt}
                onChange={(e) => setExtractedPrompt(e.target.value)}
                className="w-full flex-1 min-h-[200px] bg-zinc-950 border border-zinc-800 rounded-xl p-4 text-sm text-zinc-200 focus:outline-none focus:border-amber-500 resize-none"
              />

              <div className="flex items-center justify-between mt-4">
                <button 
                  onClick={() => processStep1()} // re-run extraction
                  disabled={isLoading}
                  className="flex items-center gap-2 text-zinc-400 hover:text-white px-4 py-2 rounded-lg hover:bg-zinc-800 transition-colors disabled:opacity-50"
                >
                  <RefreshCw size={16} /> Regenerate Extraction
                </button>
                <button 
                  onClick={processStep2}
                  disabled={isLoading}
                  className="flex items-center gap-2 bg-amber-500 hover:bg-amber-600 text-black px-6 py-3 rounded-xl font-bold transition-all disabled:opacity-50"
                >
                  {isLoading ? <Loader2 className="animate-spin" size={18} /> : <ArrowRight size={18} />}
                  Generate Images (Nano Banana)
                </button>
              </div>
            </div>
          )}

          {currentStep === 3 && (
            <div className="flex-1 flex flex-col gap-4 animate-in fade-in slide-in-from-right-8 duration-300">
              <h2 className="text-xl font-bold">Concept Art generated by Nano Banana</h2>
              <p className="text-sm text-zinc-400">
                Review the 4-sided product imagery (front, back, left, right). These images will be used by Meshy's Multi-Image API to construct the final 3D model.
              </p>

              <div className="flex-1 bg-zinc-950 border border-zinc-800 rounded-xl p-4">
                {generatedImages.length > 0 ? (
                  <div className="grid grid-cols-2 gap-4 h-full">
                    {generatedImages.map((img, idx) => (
                      <div key={idx} className="relative bg-zinc-900 rounded-lg overflow-hidden border border-zinc-800 flex items-center justify-center">
                        <img src={img} alt={`View ${idx + 1}`} className="max-w-full max-h-[150px] object-contain" />
                        <div className="absolute top-2 left-2 bg-black/50 backdrop-blur-sm px-2 py-1 rounded text-[10px] font-bold text-zinc-300 uppercase">
                          {idx === 0 ? 'FRONT' : idx === 1 ? 'BACK' : idx === 2 ? 'LEFT' : 'RIGHT'}
                        </div>
                      </div>
                    ))}
                  </div>
                ) : (
                  <div className="flex h-full items-center justify-center">
                    <ImageIcon size={48} className="text-zinc-800" />
                  </div>
                )}
              </div>

              <div className="flex items-center justify-between mt-4">
                <button 
                  onClick={() => processStep2()} // re-run generation
                  disabled={isLoading}
                  className="flex items-center gap-2 text-zinc-400 hover:text-white px-4 py-2 rounded-lg hover:bg-zinc-800 transition-colors disabled:opacity-50"
                >
                  <RefreshCw size={16} /> Regenerate Images
                </button>
                <button 
                  onClick={processStep3}
                  disabled={isLoading}
                  className="flex items-center gap-2 bg-amber-500 hover:bg-amber-600 text-black px-6 py-3 rounded-xl font-bold transition-all disabled:opacity-50"
                >
                  {isLoading ? <Loader2 className="animate-spin" size={18} /> : <ArrowRight size={18} />}
                  Construct 3D Model (Meshy)
                </button>
              </div>
            </div>
          )}

          {currentStep === 4 && (
            <div className="flex-1 flex flex-col items-center justify-center gap-6 animate-in fade-in slide-in-from-right-8 duration-300 w-full">
              {modelProgress < 100 ? (
                <>
                  <div className="w-24 h-24 rounded-full border-4 border-zinc-800 border-t-amber-500 animate-spin flex items-center justify-center">
                    <Box size={32} className="text-zinc-500 animate-pulse" />
                  </div>
                  
                  <div className="text-center w-full max-w-sm">
                    <h2 className="text-xl font-bold mb-2">Generating 3D Model...</h2>
                    <p className="text-sm text-zinc-400 mb-6">
                      Meshy is constructing the model based on the concept images.
                    </p>

                    <div className="w-full bg-zinc-950 rounded-full h-3 overflow-hidden">
                      <div 
                        className="h-full bg-amber-500 transition-all duration-300"
                        style={{ width: `${modelProgress}%` }}
                      />
                    </div>
                    <div className="text-xs font-bold text-amber-500 mt-2">{modelProgress}% Complete</div>
                  </div>
                </>
              ) : (
                <div className="w-full h-[400px] bg-zinc-900 rounded-xl overflow-hidden relative border border-zinc-800 flex flex-col items-center justify-center">
                  {generatedGlbUrl ? (
                    <Canvas shadows dpr={[1, 2]} camera={{ position: [0, 0, 4], fov: 45 }}>
                      <Suspense fallback={null}>
                        <Stage environment="studio" intensity={0.5}>
                          <ModelPreview url={generatedGlbUrl} />
                        </Stage>
                      </Suspense>
                      <OrbitControls autoRotate />
                    </Canvas>
                  ) : (
                    <div className="text-zinc-500">Failed to load preview</div>
                  )}
                  <div className="absolute bottom-6 flex justify-center w-full z-10 pointer-events-none">
                    <button 
                      onClick={processStep4}
                      disabled={isLoading}
                      className="pointer-events-auto flex items-center gap-2 bg-amber-500 hover:bg-amber-600 text-black px-8 py-3 rounded-xl font-bold transition-all shadow-xl disabled:opacity-50 animate-in fade-in slide-in-from-bottom-4"
                    >
                      {isLoading ? <Loader2 className="animate-spin" size={18} /> : <Check size={18} />}
                      Open in Editor
                    </button>
                  </div>
                </div>
              )}
            </div>
          )}

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