'use client';

import { useState, useEffect, useRef } from 'react';
import { useSearchParams, useRouter } from 'next/navigation';
import { toast } from 'react-hot-toast';
import {
    Wand2, Download, History, RotateCcw, Image as ImageIcon,
    Sparkles, Zap, ArrowRight, Loader2, Save, FolderOpen
} from 'lucide-react';
import api, { modulesApi, assetsApi } from '@/lib/api';
import { useStore } from '@/lib/store';
import FileUpload from '@/components/FileUpload';
import AssetPickerModal from '@/components/AssetPickerModal';
import JobProgress from '@/components/JobProgress';

export default function NanoBananaProPage() {
    const router = useRouter();
    const searchParams = useSearchParams();
    const initialAssetId = searchParams.get('assetId');

    // State
    const [currentAsset, setCurrentAsset] = useState<any>(null);
    const [prompt, setPrompt] = useState('');
    const [loading, setLoading] = useState(false);
    const [history, setHistory] = useState<any[]>([]);
    const [imageSize, setImageSize] = useState('2K');
    const [aspectRatio, setAspectRatio] = useState('16:9');

    // Picker State
    const [showPicker, setShowPicker] = useState(false);

    const handleAssetSelect = (assets: any[]) => {
        if (assets.length > 0) {
            const selected = assets[0];
            setCurrentAsset(selected);
            setHistory((prev: any[]) => [...prev, selected]);
            toast.success('Image selected from library');
        }
    };

    // Quick Actions (from PHP reference)
    const quickActions = [
        { label: 'Add Lighting', prompt: 'Add dramatic lighting' },
        { label: 'Add Sunset', prompt: 'Add sunset in background' },
        { label: 'More Vibrant', prompt: 'Make colors more vibrant' },
        { label: 'Motion Blur', prompt: 'Add motion blur' },
        { label: 'Photorealistic', prompt: 'Make it photorealistic' },
        { label: 'Depth of Field', prompt: 'Add depth of field effect' },
    ];

    // Load initial asset if provided
    useEffect(() => {
        if (initialAssetId) {
            loadAsset(initialAssetId);
        }
    }, [initialAssetId]);

    const loadAsset = async (id: string) => {
        try {
            const response = await assetsApi.get(id);
            setCurrentAsset(response.data);
            // Add to history if not present
            setHistory(prev => {
                if (!prev.find(item => item.id === id)) {
                    return [...prev, response.data];
                }
                return prev;
            });
        } catch (err) {
            toast.error('Failed to load initial image');
        }
    };

    const handleGenerate = async (e?: React.FormEvent) => {
        if (e) e.preventDefault();
        if (!prompt.trim()) {
            toast.error('Please enter a prompt');
            return;
        }

        setLoading(true);
        try {
            // Logic for Edit Loop:
            // The backend endpoint /modules/image/generate expects a JSON payload (ImageGenerateRequest).
            // It supports 'reference_asset_id' to load the image from the server side.
            // We do NOT upload the file here; currentAsset is already on the server.

            const payload: any = {
                prompt: prompt,
                provider: 'nano-banana',
                model: 'gemini-3-pro-image-preview', // Enforce Nano Banana
                // Standard params
                aspect_ratio: aspectRatio,
                provider_options: {
                    image_size: imageSize,
                }
            };

            if (currentAsset) {
                payload.reference_asset_id = currentAsset.id;
                console.log('✓ Adding reference asset:', currentAsset.id, currentAsset.original_filename);
            } else {
                console.warn('⚠️ No currentAsset - generating from scratch');
            }

            console.log('📤 Sending payload:', JSON.stringify(payload, null, 2));

            // Use direct API call with JSON
            const response = await api.post('/modules/image/generate', payload);

            // The response is a Job ID
            const jobId = response.data.id;

            // Set active job ID to show progress bar
            setActiveJobId(jobId);

        } catch (err: any) {
            console.error(err);
            toast.error(err.response?.data?.detail || 'Generation failed');
            setLoading(false);
        }
    };

    // (pollJob removed as we use JobProgress component)

    // Revised Strategy: 
    // We'll use a `JobProgress` component that calls a callback `onComplete`
    // The `handleGenerate` will just kick off the job and set a local `activeJobId`.

    const [activeJobId, setActiveJobId] = useState<string | null>(null);

    const handleJobComplete = async (job: any) => {
        setLoading(false);
        setActiveJobId(null);
        if (job.status === 'completed' && job.output_asset_ids?.length > 0) {
            // Fetch the new asset
            const newAssetId = job.output_asset_ids[0];
            try {
                const assetResponse = await assetsApi.get(newAssetId);
                const newAsset = assetResponse.data;
                setCurrentAsset(newAsset);
                setHistory(prev => [...prev, newAsset]);
                setPrompt(''); // Clear prompt on success for next edit
                toast.success('Image updated!');
            } catch (e) {
                toast.error('Failed to load new image');
            }
        } else {
            toast.error('Generation failed');
        }
    };

    const handleJobError = (msg: string) => {
        setLoading(false);
        setActiveJobId(null);
        toast.error(msg);
    };

    const startJob = async (jobId: string) => {
        setActiveJobId(jobId);
    };

    // Quick Action Handler
    const handleQuickAction = (actionPrompt: string) => {
        setPrompt(actionPrompt);
        // Auto-submit? The PHP code does.
        // Let's set it and let user confirm or auto-submit.
        // User requested "works on tha nano banan generate and edit cycle", implies speed.
        // Let's just set the prompt and scroll to text area.
    };

    const handleReset = () => {
        if (confirm('Start fresh? This will clear the current image context.')) {
            setCurrentAsset(null);
            setPrompt('');
            router.push('/image/edit-pro'); // Clear URL param
        }
    };

    const handleFileUpload = async (file: File) => {
        setLoading(true);
        try {
            const response = await assetsApi.upload(file, undefined, true);
            const uploadedAsset = response.data;
            setCurrentAsset(uploadedAsset);
            setHistory(prev => [...prev, uploadedAsset]);
            toast.success('Image uploaded ready for editing');
        } catch (err) {
            toast.error('Failed to upload image');
        } finally {
            setLoading(false);
        }
    };

    const handleDownload = async () => {
        if (!currentAsset) return;
        try {
            const response = await assetsApi.download(currentAsset.id);
            const url = window.URL.createObjectURL(response.data);
            const a = document.createElement('a');
            a.href = url;
            a.download = currentAsset.original_filename;
            a.click();
            window.URL.revokeObjectURL(url);
        } catch (e) {
            toast.error('Download failed');
        }
    };

    return (
        <div className="min-h-screen bg-black text-white p-6 md:p-12 font-montserrat">
            <div className="max-w-[1600px] mx-auto grid grid-cols-1 lg:grid-cols-12 gap-8">

                {/* Header */}
                <div className="lg:col-span-12 flex items-center justify-between border-b-2 border-forge-yellow pb-6 mb-4">
                    <div>
                        <h1 className="text-4xl font-bold text-forge-yellow tracking-tight mb-2">Nano Banana Pro</h1>
                        <p className="text-gray-400">Advanced AI Image Generation & Iterative Editing Loop</p>
                    </div>
                    <div className="flex gap-2">
                        <button onClick={handleReset} className="px-4 py-2 rounded-lg bg-gray-800 hover:bg-gray-700 text-gray-300 transition-colors flex items-center gap-2">
                            <RotateCcw className="w-4 h-4" /> Reset
                        </button>
                    </div>
                </div>

                {/* Left Panel: Controls */}
                <div className="lg:col-span-4 space-y-6">
                    <div className="bg-[#1a1a1a] border border-[#333] rounded-xl p-6">
                        <h2 className="text-xl font-semibold text-forge-yellow mb-6 flex items-center gap-2">
                            <Wand2 className="w-5 h-5" />
                            {currentAsset ? 'Edit Instruction' : 'Create New Image'}
                        </h2>

                        <div className="space-y-4">
                            <div>
                                <div className="flex items-center justify-between mb-2">
                                    <label className="block text-forge-yellow font-medium text-sm">Upload Initial Image (Optional)</label>
                                    <button
                                        onClick={() => setShowPicker(true)}
                                        className="text-xs flex items-center gap-1 text-gray-400 hover:text-white transition-colors"
                                    >
                                        <FolderOpen className="w-3 h-3" /> Select from Library
                                    </button>
                                </div>
                                <FileUpload
                                    onUpload={handleFileUpload}
                                    accept={{ 'image/*': ['.png', '.jpg', '.jpeg', '.webp'] }}
                                    currentFile={null}
                                    label="Drop an image to start editing"
                                    className="bg-black border-2 border-[#333] hover:border-forge-yellow transition-colors"
                                />

                                <AssetPickerModal
                                    isOpen={showPicker}
                                    onClose={() => setShowPicker(false)}
                                    onConfirm={handleAssetSelect}
                                    allowedTypes={['image/']}
                                    maxSelect={1}
                                    title="Select Image to Edit"
                                />

                            </div>

                            <div>
                                <label className="block text-forge-yellow font-medium mb-2 text-sm">Prompt</label>
                                <textarea
                                    value={prompt}
                                    onChange={(e) => setPrompt(e.target.value)}
                                    placeholder={currentAsset ? 'e.g., "Add a sunset tone", "remove the person", "make it oil painting"' : 'Describe the image you want to create...'}
                                    className="w-full h-32 bg-black border-2 border-[#333] focus:border-forge-yellow rounded-lg p-4 text-white placeholder-gray-600 outline-none resize-none transition-colors"
                                    onKeyDown={(e) => {
                                        if (e.key === 'Enter' && !e.shiftKey) {
                                            e.preventDefault();
                                            handleGenerate();
                                        }
                                    }}
                                />
                            </div>

                            {!currentAsset && (
                                <div className="grid grid-cols-2 gap-4">
                                    <div>
                                        <label className="block text-forge-yellow font-medium mb-2 text-sm">Aspect Ratio</label>
                                        <select
                                            value={aspectRatio}
                                            onChange={(e) => setAspectRatio(e.target.value)}
                                            className="w-full bg-black border border-[#333] rounded-lg p-2 outline-none focus:border-forge-yellow"
                                        >
                                            <option value="16:9">16:9 (Landscape)</option>
                                            <option value="1:1">1:1 (Square)</option>
                                            <option value="9:16">9:16 (Portrait)</option>
                                            <option value="4:3">4:3</option>
                                            <option value="3:4">3:4</option>
                                        </select>
                                    </div>
                                    <div>
                                        <label className="block text-forge-yellow font-medium mb-2 text-sm">Size</label>
                                        <select
                                            value={imageSize}
                                            onChange={(e) => setImageSize(e.target.value)}
                                            className="w-full bg-black border border-[#333] rounded-lg p-2 outline-none focus:border-forge-yellow"
                                        >
                                            <option value="1K">1K (Fast)</option>
                                            <option value="2K">2K (Balanced)</option>
                                        </select>
                                    </div>
                                </div>
                            )}

                            <button
                                onClick={(e) => handleGenerate(e)}
                                disabled={loading || !prompt.trim()}
                                className={`w-full py-4 rounded-lg font-bold text-lg transition-all transform active:scale-95 ${loading || !prompt.trim()
                                    ? 'bg-gray-700 text-gray-400 cursor-not-allowed'
                                    : 'bg-forge-yellow text-black hover:bg-[#ffcd1f] hover:-translate-y-1 shadow-[0_4px_12px_rgba(255,196,7,0.2)]'
                                    }`}
                            >
                                {loading ? (
                                    <div className="flex items-center justify-center gap-2">
                                        <Loader2 className="w-5 h-5 animate-spin" />
                                        {currentAsset ? 'Editing...' : 'Generating...'}
                                    </div>
                                ) : (
                                    <div className="flex items-center justify-center gap-2">
                                        {currentAsset ? <Sparkles className="w-5 h-5" /> : <Zap className="w-5 h-5" />}
                                        {currentAsset ? 'Update Image' : 'Generate'}
                                    </div>
                                )}
                            </button>
                        </div>

                        {/* Quick Actions */}
                        {currentAsset && (
                            <div className="mt-8 pt-6 border-t border-[#333]">
                                <h3 className="text-gray-400 text-xs font-bold uppercase tracking-wider mb-4">Quick Edits</h3>
                                <div className="grid grid-cols-2 gap-2">
                                    {quickActions.map((action, idx) => (
                                        <button
                                            key={idx}
                                            onClick={() => handleQuickAction(action.prompt)}
                                            className="p-2 text-xs bg-[#222] text-forge-yellow border border-forge-yellow/30 hover:bg-forge-yellow hover:text-black rounded transition-all font-medium text-center"
                                        >
                                            {action.label}
                                        </button>
                                    ))}
                                </div>
                            </div>
                        )}

                        {/* Active Job Tracker */}
                        {activeJobId && (
                            <div className="mt-6">
                                <JobProgress jobId={activeJobId} onComplete={handleJobComplete} onError={handleJobError} />
                            </div>
                        )}

                    </div>

                    {/* History List (Simple) */}
                    {history.length > 1 && (
                        <div className="bg-[#1a1a1a] border border-[#333] rounded-xl p-4">
                            <h3 className="text-gray-400 text-xs font-bold uppercase tracking-wider mb-4 flex items-center gap-2">
                                <History className="w-4 h-4" /> Session History
                            </h3>
                            <div className="space-y-2 max-h-48 overflow-y-auto custom-scrollbar">
                                {history.slice(0, -1).reverse().map((histAsset, index) => (
                                    <div
                                        key={`${histAsset.id}-${index}`}
                                        onClick={() => setCurrentAsset(histAsset)}
                                        className="flex items-center gap-3 p-2 rounded hover:bg-[#2a2a2a] cursor-pointer group"
                                    >
                                        <div className="w-12 h-12 rounded bg-black overflow-hidden shrink-0 border border-gray-700">
                                            <img
                                                src={`/api/v1/assets/${histAsset.id}/download`}
                                                className="w-full h-full object-cover opacity-70 group-hover:opacity-100 transition-opacity"
                                                alt="History"
                                            />
                                        </div>
                                        <div className="min-w-0">
                                            <p className="text-xs text-gray-300 truncate">{histAsset.created_at}</p>
                                            <p className="text-[10px] text-gray-500 truncate">ID: {histAsset.id.slice(0, 8)}</p>
                                        </div>
                                    </div>
                                ))}
                            </div>
                        </div>
                    )}
                </div>

                {/* Right Panel: Display */}
                <div className="lg:col-span-8 bg-[#0a0a0a] rounded-2xl border-2 border-[#1a1a1a] p-2 flex flex-col h-[800px]">
                    <div className="flex-1 flex items-center justify-center bg-black/50 rounded-xl relative overflow-hidden group">
                        {currentAsset ? (
                            <>
                                <img
                                    src={`/api/v1/assets/${currentAsset.id}/download`}
                                    className="max-w-full max-h-full object-contain shadow-2xl"
                                    alt="Current"
                                />
                                <div className="absolute top-4 right-4 flex gap-2 opacity-0 group-hover:opacity-100 transition-opacity">
                                    <button
                                        onClick={handleDownload}
                                        className="bg-black/80 hover:bg-forge-yellow hover:text-black text-white p-3 rounded-full backdrop-blur-md transition-all border border-white/10"
                                    >
                                        <Download className="w-5 h-5" />
                                    </button>
                                </div>
                            </>
                        ) : (
                            <div className="text-center text-[#333] space-y-4">
                                <div className="w-24 h-24 rounded-full border-4 border-[#1a1a1a] flex items-center justify-center mx-auto">
                                    <ImageIcon className="w-10 h-10" />
                                </div>
                                <p className="font-medium text-lg">Ready to Create</p>
                                <p className="text-sm max-w-xs mx-auto">Enter a prompt to generate your first image, then iterate with precise edits.</p>
                            </div>
                        )}
                    </div>
                </div>

            </div>
        </div>
    );
}

// End of file
