import React, { useState, useEffect } from 'react';
import { motion, AnimatePresence } from 'framer-motion';
import { X, Sparkles, Image as ImageIcon, Loader2, Download, Settings, Wand2, Lightbulb } from 'lucide-react';
import { cn } from '../lib/utils';
import { useAuthStore } from '../store/authStore';
import { api } from '../lib/api';

interface VizAiModalProps {
  isOpen: boolean;
  onClose: () => void;
  baseImage: string | null;
}

const dataURLtoBlob = (dataurl: string) => {
  const arr = dataurl.split(',');
  const mime = arr[0].match(/:(.*?);/)?.[1];
  const bstr = atob(arr[1]);
  let n = bstr.length;
  const u8arr = new Uint8Array(n);
  while(n--) {
      u8arr[n] = bstr.charCodeAt(n);
  }
  return new Blob([u8arr], {type:mime});
};

export const VizAiModal: React.FC<VizAiModalProps> = ({ isOpen, onClose, baseImage }) => {
  const { profile, setProfile } = useAuthStore();
  
  // API Keys
  const [stabilityApiKey, setStabilityApiKey] = useState('');
  const [openAiKey, setOpenAiKey] = useState('');
  
  // Prompt Generation State (AI Idea)
  const [productName, setProductName] = useState('');
  const [ingredients, setIngredients] = useState('');
  const [imageTone, setImageTone] = useState('');
  const [packDesign, setPackDesign] = useState('');
  const [isGeneratingPrompt, setIsGeneratingPrompt] = useState(false);
  
  // Image Generation State
  const [prompt, setPrompt] = useState('');
  const [strength, setStrength] = useState(0.85);
  const [isGeneratingImage, setIsGeneratingImage] = useState(false);
  const [generatedImage, setGeneratedImage] = useState<string | null>(null);
  const [error, setError] = useState<string | null>(null);

  useEffect(() => {
    if (isOpen) {
      setStabilityApiKey(localStorage.getItem('stability_api_key') || '');
      setOpenAiKey(localStorage.getItem('openai_api_key') || '');
    }
  }, [isOpen]);

  const generateAiIdea = async () => {
    if (!openAiKey) {
      setError("Please set your OpenAI API key in Tools > API Settings first.");
      return;
    }
    
    setIsGeneratingPrompt(true);
    setError(null);
    try {
      const response = await fetch('https://api.openai.com/v1/chat/completions', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${openAiKey}`
        },
        body: JSON.stringify({
          model: 'gpt-4o-mini',
          messages: [{
            role: 'system',
            content: 'You are an expert prompt engineer for Stable Diffusion. Your goal is to take product specifications and write a single, highly descriptive prompt for generating a beautiful, professional product mockup scene. Describe the lighting, the background, and the aesthetic vividly.'
          }, {
            role: 'user',
            content: `Create a single, highly descriptive image generation prompt based on these details:
            Product Name: ${productName || 'Unspecified Product'}
            Ingredients/Elements: ${ingredients || 'None specified'}
            Image Tone/Mood: ${imageTone || 'Professional, studio lighting'}
            Pack Design/Shape: ${packDesign || 'Standard'}
            
            Return ONLY the prompt string (under 100 words) without any conversational text or quotes.`
          }]
        })
      });
      
      const data = await response.json();
      if (data.error) throw new Error(data.error.message);
      
      const generated = data.choices?.[0]?.message?.content?.trim() || "";
      // Strip quotes if they were included
      setPrompt(generated.replace(/^"|"$/g, ''));
    } catch (err: any) {
      console.error(err);
      setError(err.message || "Failed to generate AI Idea.");
    } finally {
      setIsGeneratingPrompt(false);
    }
  };

  const generateImage = async () => {
    if (!stabilityApiKey) {
      setError("Please set your Stability AI API key in Tools > API Settings first.");
      return;
    }
    if (!baseImage) {
      setError("No base image captured from the scene.");
      return;
    }
    if (!prompt) {
      setError("Please enter a prompt or use AI Idea.");
      return;
    }
    
    if (!profile || profile.api_credits <= 0) {
      setError("Not enough AI Credits. Please contact an administrator.");
      return;
    }

    setIsGeneratingImage(true);
    setError(null);

    try {
      const formData = new FormData();
      formData.append('image', dataURLtoBlob(baseImage), 'base.png');
      formData.append('prompt', prompt);
      formData.append('strength', strength.toString());
      formData.append('mode', 'image-to-image');
      formData.append('model', 'sd3.5-large');
      formData.append('output_format', 'png');

      const response = await fetch('https://api.stability.ai/v2beta/stable-image/generate/sd3', {
        method: 'POST',
        headers: {
          'Authorization': `Bearer ${stabilityApiKey}`,
          'Accept': 'image/*'
        },
        body: formData
      });

      if (!response.ok) {
        const errData = await response.json().catch(() => ({}));
        throw new Error(errData.message || `Stability API error: ${response.status}`);
      }

      const blob = await response.blob();
      const imageUrl = URL.createObjectURL(blob);
      setGeneratedImage(imageUrl);
      
      // Deduct credit
      if (profile) {
        const { api_credits } = await api.deductCredit();
        setProfile({ ...profile, api_credits });
      }
      
    } catch (err: any) {
      console.error(err);
      setError(err.message || "Failed to generate image.");
    } finally {
      setIsGeneratingImage(false);
    }
  };

  if (!isOpen) return null;

  return (
    <AnimatePresence>
      <div className="fixed inset-0 z-[100] flex items-center justify-center p-4 bg-black/60 backdrop-blur-sm">
        <motion.div 
          initial={{ opacity: 0, scale: 0.95, y: 20 }}
          animate={{ opacity: 1, scale: 1, y: 0 }}
          exit={{ opacity: 0, scale: 0.95, y: 20 }}
          className="w-full max-w-5xl bg-zinc-900 border border-zinc-800 rounded-xl shadow-2xl overflow-hidden flex flex-col h-[80vh]"
        >
          {/* Header */}
          <div className="h-14 bg-zinc-950 flex items-center justify-between px-6 border-b border-zinc-800 shrink-0">
            <div className="flex items-center gap-3">
              <div className="w-8 h-8 rounded-lg bg-gradient-to-tr from-purple-600 to-blue-500 flex items-center justify-center">
                <Sparkles size={16} className="text-white" />
              </div>
              <div>
                <h2 className="text-sm font-bold text-zinc-100">Titan3d AI Generator</h2>
                <p className="text-[10px] text-zinc-400">Generate stunning product mockups directly from your 3D scene</p>
              </div>
            </div>
            
            <div className="flex items-center gap-2">
              <button onClick={onClose} className="p-2 hover:bg-zinc-800 rounded-lg transition-colors">
                <X size={18} className="text-zinc-400" />
              </button>
            </div>
          </div>

          <div className="flex-1 overflow-hidden flex bg-zinc-900/50">
              <div className="flex w-full h-full">
                {/* Left Panel: Inputs */}
                <div className="w-[400px] border-r border-zinc-800 flex flex-col overflow-y-auto custom-scrollbar">
                  {/* AI Idea Box */}
                  <div className="p-5 border-b border-zinc-800 bg-zinc-950/30">
                    <div className="flex items-center gap-2 mb-4">
                      <Lightbulb size={16} className="text-yellow-500" />
                      <h3 className="text-xs font-bold text-zinc-300 uppercase tracking-wider">AI Idea Generator</h3>
                    </div>
                    
                    <div className="flex flex-col gap-3">
                      <input 
                        type="text" 
                        placeholder="Product Name (e.g. Tetra Juice)" 
                        value={productName}
                        onChange={(e) => setProductName(e.target.value)}
                        className="w-full bg-zinc-900 border border-zinc-800 rounded px-3 py-2 text-xs text-white placeholder-zinc-600 focus:border-yellow-500/50 outline-none"
                      />
                      <input 
                        type="text" 
                        placeholder="Ingredients (e.g. Oranges, Mint)" 
                        value={ingredients}
                        onChange={(e) => setIngredients(e.target.value)}
                        className="w-full bg-zinc-900 border border-zinc-800 rounded px-3 py-2 text-xs text-white placeholder-zinc-600 focus:border-yellow-500/50 outline-none"
                      />
                      <input 
                        type="text" 
                        placeholder="Image Tone (e.g. Fresh, Bright, Morning)" 
                        value={imageTone}
                        onChange={(e) => setImageTone(e.target.value)}
                        className="w-full bg-zinc-900 border border-zinc-800 rounded px-3 py-2 text-xs text-white placeholder-zinc-600 focus:border-yellow-500/50 outline-none"
                      />
                      <input 
                        type="text" 
                        placeholder="Pack Design (e.g. Sleek, Modern)" 
                        value={packDesign}
                        onChange={(e) => setPackDesign(e.target.value)}
                        className="w-full bg-zinc-900 border border-zinc-800 rounded px-3 py-2 text-xs text-white placeholder-zinc-600 focus:border-yellow-500/50 outline-none"
                      />
                      <button 
                        onClick={generateAiIdea}
                        disabled={isGeneratingPrompt}
                        className="w-full mt-2 py-2 bg-yellow-500/10 hover:bg-yellow-500/20 text-yellow-500 border border-yellow-500/30 rounded-lg text-xs font-bold transition-all flex items-center justify-center gap-2"
                      >
                        {isGeneratingPrompt ? <Loader2 size={14} className="animate-spin" /> : <Wand2 size={14} />}
                        {isGeneratingPrompt ? 'Generating Idea...' : 'Suggest Prompt'}
                      </button>
                    </div>
                  </div>

                  {/* Main Generation Settings */}
                  <div className="p-5 flex flex-col gap-6">
                    <div className="flex flex-col gap-2">
                      <label className="text-xs font-bold text-zinc-300">Prompt</label>
                      <textarea 
                        value={prompt}
                        onChange={(e) => setPrompt(e.target.value)}
                        placeholder="Describe the background and lighting..."
                        className="w-full h-32 bg-zinc-900 border border-zinc-800 rounded-lg p-3 text-sm text-white placeholder-zinc-600 focus:border-blue-500/50 outline-none resize-none"
                      />
                    </div>

                    <div className="flex flex-col gap-2">
                      <div className="flex justify-between">
                        <label className="text-xs font-bold text-zinc-400">Creative Strength</label>
                        <span className="text-xs text-blue-400 font-mono">{strength.toFixed(2)}</span>
                      </div>
                      <input 
                        type="range" 
                        min="0" max="1" step="0.05"
                        value={strength}
                        onChange={(e) => setStrength(parseFloat(e.target.value))}
                        className="w-full accent-blue-500"
                      />
                      <div className="flex justify-between text-[10px] text-zinc-600">
                        <span>Similar to Base</span>
                        <span>Highly Creative</span>
                      </div>
                    </div>

                    {error && (
                      <div className="p-3 bg-red-500/10 border border-red-500/20 rounded-lg text-xs text-red-400">
                        {error}
                      </div>
                    )}

                    <button 
                      onClick={generateImage}
                      disabled={isGeneratingImage}
                      className="w-full py-4 mt-auto bg-gradient-to-r from-purple-600 to-blue-600 hover:from-purple-500 hover:to-blue-500 text-white font-bold rounded-xl shadow-lg transition-all flex items-center justify-center gap-2"
                    >
                      {isGeneratingImage ? <Loader2 size={18} className="animate-spin" /> : <Sparkles size={18} />}
                      {isGeneratingImage ? 'Generating Scene...' : 'Generate Image'}
                    </button>
                  </div>
                </div>

                {/* Right Panel: Previews */}
                <div className="flex-1 p-6 bg-zinc-950 flex flex-col gap-4 relative overflow-y-auto">
                  <div className="grid grid-cols-2 gap-6 h-full">
                    {/* Base Image */}
                    <div className="flex flex-col gap-2">
                      <div className="flex items-center justify-between">
                        <span className="text-xs font-bold text-zinc-500 uppercase tracking-widest">Base Render</span>
                      </div>
                      <div className="flex-1 bg-zinc-900 border border-zinc-800 rounded-xl overflow-hidden flex items-center justify-center relative shadow-inner">
                        {baseImage ? (
                          <img src={baseImage} alt="Base Render" className="w-full h-full object-contain" />
                        ) : (
                          <div className="text-zinc-600 flex flex-col items-center gap-2">
                            <ImageIcon size={32} />
                            <span className="text-xs">No scene captured</span>
                          </div>
                        )}
                      </div>
                    </div>

                    {/* Generated Image */}
                    <div className="flex flex-col gap-2">
                      <div className="flex items-center justify-between">
                        <span className="text-xs font-bold text-blue-500 uppercase tracking-widest">AI Result</span>
                        {generatedImage && (
                          <a 
                            href={generatedImage} 
                            download="viz_ai_mockup.png"
                            className="text-xs text-zinc-400 hover:text-white flex items-center gap-1"
                          >
                            <Download size={14} /> Download
                          </a>
                        )}
                      </div>
                      <div className="flex-1 bg-zinc-900 border border-blue-500/20 rounded-xl overflow-hidden flex items-center justify-center relative shadow-lg">
                        {isGeneratingImage ? (
                          <div className="flex flex-col items-center gap-4 text-blue-500">
                            <Loader2 size={32} className="animate-spin" />
                            <span className="text-sm font-medium animate-pulse">Rendering via Stability AI...</span>
                          </div>
                        ) : generatedImage ? (
                          <img src={generatedImage} alt="AI Result" className="w-full h-full object-contain" />
                        ) : (
                          <div className="text-zinc-600 flex flex-col items-center gap-2">
                            <Sparkles size={32} className="opacity-50" />
                            <span className="text-xs">Result will appear here</span>
                          </div>
                        )}
                      </div>
                    </div>
                  </div>
                </div>
              </div>
          </div>
        </motion.div>
      </div>
    </AnimatePresence>
  );
};