'use client';

import { useState, useEffect } from 'react';
import { useRouter, useSearchParams } from 'next/navigation';
import { toast } from 'react-hot-toast';
import { ImagePlus, Download, Sparkles, Pencil, X, Loader2, Maximize, Film } from 'lucide-react';
import JobProgress from '@/components/JobProgress';
import ProviderControls from '@/components/ProviderControls';
import { modulesApi, assetsApi, capabilitiesApi } from '@/lib/api';
import { useStore } from '@/lib/store';
import { ProviderConfig } from '@/types/providers';

export default function ImageGeneratePage() {
  const { addJob, updateJob } = useStore();
  const router = useRouter();
  const searchParams = useSearchParams();

  // Provider config state
  const [capabilities, setCapabilities] = useState<Record<string, ProviderConfig> | null>(null);
  const [loadingCapabilities, setLoadingCapabilities] = useState(true);

  // Generation state
  const [prompt, setPrompt] = useState('');
  const [provider, setProvider] = useState('openai');
  const [model, setModel] = useState('');
  const [providerOptions, setProviderOptions] = useState<Record<string, any>>({});
  const [jobId, setJobId] = useState<string | null>(null);
  const [generatedImages, setGeneratedImages] = useState<any[]>([]);
  const [loading, setLoading] = useState(false);

  // Iterative editing state (for Nano Banana)
  const [editingImage, setEditingImage] = useState<any | null>(null);
  const [editInstructions, setEditInstructions] = useState('');

  // Load provider capabilities on mount
  useEffect(() => {
    const loadCapabilities = async () => {
      try {
        const response = await capabilitiesApi.getImageProviders();
        setCapabilities(response.data);

        // Set default provider and model
        const firstProvider = Object.keys(response.data)[0];
        setProvider(firstProvider);
        setModel(response.data[firstProvider].defaultModel);

        // Initialize with default values
        initializeDefaults(response.data[firstProvider]);
      } catch (err) {
        console.error('Failed to load provider configurations:', err);
        toast.error('Failed to load provider configurations');
      } finally {
        setLoadingCapabilities(false);
      }
    };

    loadCapabilities();
  }, []);

  // Handle URL parameters
  useEffect(() => {
    const urlPrompt = searchParams.get('prompt');
    if (urlPrompt) {
      setPrompt(urlPrompt);
    }

    // Check for reference asset for editing/variations if we support it via URL
    // (Optional: handle assetId if needed)
  }, [searchParams]);

  // Initialize default values for provider
  const initializeDefaults = (config: ProviderConfig) => {
    if (!config) {
      console.error('Config is undefined');
      return;
    }

    const defaults: Record<string, any> = {};

    // Common controls
    if (config.commonControls && Array.isArray(config.commonControls)) {
      config.commonControls.forEach((control) => {
        defaults[control.name] = control.default;
      });
    }

    // Model-specific controls
    const modelConfig = config.models?.find(m => m.id === config.defaultModel);
    if (modelConfig?.controls && Array.isArray(modelConfig.controls)) {
      modelConfig.controls.forEach((control) => {
        defaults[control.name] = control.default;
      });
    }

    setProviderOptions(defaults);
  };

  // Handle provider change
  const handleProviderChange = (newProvider: string) => {
    if (!capabilities) return;

    setProvider(newProvider);
    const config = capabilities[newProvider];
    setModel(config.defaultModel);
    initializeDefaults(config);

    // Cancel editing if changing provider
    if (editingImage) {
      setEditingImage(null);
      setEditInstructions('');
    }
  };

  // Handle model change
  const handleModelChange = (newModel: string) => {
    setModel(newModel);

    if (!capabilities) return;
    const config = capabilities[provider];
    const modelConfig = config.models.find(m => m.id === newModel);

    // Merge current options with model defaults
    const modelDefaults: Record<string, any> = {};
    modelConfig?.controls?.forEach((control) => {
      if (!(control.name in providerOptions)) {
        modelDefaults[control.name] = control.default;
      }
    });

    setProviderOptions({
      ...providerOptions,
      ...modelDefaults
    });
  };

  const handleGenerate = async () => {
    const effectivePrompt = editingImage ? editInstructions : prompt;
    if (!effectivePrompt.trim()) {
      toast.error(editingImage ? 'Please enter edit instructions' : 'Please enter a prompt');
      return;
    }

    setLoading(true);
    setGeneratedImages([]); // Always clear previous images
    setJobId(null); // Reset job ID

    try {
      const payload = {
        prompt: effectivePrompt,
        provider: editingImage ? 'nano-banana' : provider,
        model: editingImage ? 'gemini-2.5-flash-image' : model,
        provider_options: editingImage ? undefined : providerOptions,
        reference_asset_id: editingImage?.id || undefined,
      };

      // Debug logging
      if (editingImage) {
        console.log('🎨 Nano Banana Edit Mode:');
        console.log('  Reference Asset ID:', editingImage.id);
        console.log('  Edit Instructions:', effectivePrompt);
        console.log('  Full Payload:', payload);
      }

      const response = await modulesApi.generateImage(payload);

      const job = response.data;
      setJobId(job.id);
      addJob({
        id: job.id,
        module: 'image_generation',
        status: job.status,
        progress: job.progress,
        created_at: job.created_at,
      });

      toast.success(editingImage ? 'Image editing started!' : 'Image generation started!');
    } catch (err: any) {
      console.error('Generation error:', err);
      toast.error(err.response?.data?.detail || 'Failed to start generation');
      setLoading(false);
    }
  };

  const handleJobComplete = async (job: any) => {
    setLoading(false);
    updateJob(job.id, { status: 'completed', progress: 100 });

    if (job.output_asset_ids?.length > 0) {
      const images = await Promise.all(
        job.output_asset_ids.map(async (id: string) => {
          const asset = await assetsApi.get(id);
          return asset.data;
        })
      );

      // When editing, replace with new edited version to maintain chain visibility
      if (editingImage) {
        setGeneratedImages(images);
        // Keep the editing image as the new base for next edit
        setEditingImage(images[0]);
        setEditInstructions('');
        toast.success('Image edited successfully! You can continue editing.');
      } else {
        setGeneratedImages(images);
        toast.success('Images generated successfully!');
      }
    }
  };

  const handleStartEdit = (image: any) => {
    setEditingImage(image);
    setEditInstructions('');
    // Auto-switch to Nano Banana for editing
    if (capabilities && capabilities['nano-banana']) {
      setProvider('nano-banana');
      setModel('gemini-2.5-flash-image');
      initializeDefaults(capabilities['nano-banana']);
    }
  };

  const handleCancelEdit = () => {
    setEditingImage(null);
    setEditInstructions('');
  };

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

  const handleDownload = async (assetId: string, filename: string) => {
    try {
      const response = await assetsApi.download(assetId);
      const url = window.URL.createObjectURL(response.data);
      const a = document.createElement('a');
      a.href = url;
      a.download = filename;
      a.click();
      window.URL.revokeObjectURL(url);
    } catch (err) {
      toast.error('Failed to download image');
    }
  };

  if (loadingCapabilities) {
    return (
      <div className="flex items-center justify-center min-h-screen">
        <Loader2 className="w-8 h-8 animate-spin text-forge-yellow" />
      </div>
    );
  }

  const currentConfig = capabilities?.[provider];
  const supportsEditing = provider === 'nano-banana' || provider === 'gemini';

  return (
    <div className="max-w-6xl mx-auto space-y-8">
      <div className="flex items-center gap-4">
        <div className="w-12 h-12 bg-forge-yellow/10 rounded-lg flex items-center justify-center">
          <ImagePlus className="w-6 h-6 text-forge-yellow" />
        </div>
        <div>
          <h1 className="text-2xl font-bold text-white">Image Generator</h1>
          <p className="text-gray-500">Create stunning images with AI</p>
        </div>
      </div>

      <div className="grid grid-cols-1 lg:grid-cols-2 gap-8">
        {/* Controls */}
        <div className="space-y-6">
          {/* Editing Mode Panel */}
          {editingImage && (
            <div className="bg-purple-900/20 border border-purple-500/50 rounded-xl p-4 space-y-4">
              <div className="flex items-center justify-between">
                <h3 className="text-purple-400 font-medium flex items-center gap-2">
                  <Pencil className="w-4 h-4" />
                  Editing Image
                </h3>
                <button
                  onClick={handleCancelEdit}
                  className="text-gray-400 hover:text-white transition-colors"
                >
                  <X className="w-5 h-5" />
                </button>
              </div>
              <div className="flex gap-4">
                <div className="w-24 h-24 rounded-lg overflow-hidden border border-purple-500/50">
                  <img
                    src={`/api/v1/assets/${editingImage.id}/download`}
                    alt="Reference"
                    className="w-full h-full object-cover"
                  />
                </div>
                <div className="flex-1">
                  <label className="block text-sm font-medium text-gray-300 mb-2">
                    Edit Instructions
                  </label>
                  <textarea
                    value={editInstructions}
                    onChange={(e) => setEditInstructions(e.target.value)}
                    placeholder="Describe how you want to modify this image..."
                    className="input-field min-h-[80px] resize-none"
                  />
                </div>
              </div>
              <p className="text-xs text-gray-500">
                Using Nano Banana (Gemini) for iterative editing
                <br />
                Reference: {editingImage.original_filename || editingImage.id.substring(0, 8)}
              </p>
            </div>
          )}

          {/* Prompt */}
          {!editingImage && (
            <div>
              <label className="block text-sm font-medium text-gray-300 mb-2">
                Prompt
              </label>
              <textarea
                value={prompt}
                onChange={(e) => setPrompt(e.target.value)}
                placeholder="Describe the image you want to create..."
                className="input-field min-h-[120px] resize-none"
              />
            </div>
          )}

          {/* Provider & Model Selection */}
          <div className="grid grid-cols-2 gap-4">
            <div>
              <label className="block text-sm font-medium text-gray-300 mb-2">
                Provider
              </label>
              <select
                value={provider}
                onChange={(e) => handleProviderChange(e.target.value)}
                className="select-field"
                disabled={editingImage !== null}
              >
                {capabilities && Object.entries(capabilities).map(([id, config]) => (
                  <option key={id} value={id}>
                    {config.name}
                  </option>
                ))}
              </select>
            </div>
            <div>
              <label className="block text-sm font-medium text-gray-300 mb-2">
                Model
              </label>
              <select
                value={model}
                onChange={(e) => handleModelChange(e.target.value)}
                className="select-field"
                disabled={editingImage !== null}
              >
                {currentConfig?.models.map((m) => (
                  <option key={m.id} value={m.id}>
                    {m.name}
                  </option>
                ))}
              </select>
            </div>
          </div>

          {/* Model Description */}
          {currentConfig?.models.find(m => m.id === model)?.description && (
            <p className="text-xs text-gray-500 -mt-4">
              {currentConfig.models.find(m => m.id === model)?.description}
            </p>
          )}

          {/* Dynamic Provider Controls */}
          {currentConfig && !editingImage && (
            <div className="bg-forge-gray rounded-lg p-4">
              <ProviderControls
                config={currentConfig}
                selectedModel={model}
                values={providerOptions}
                onChange={setProviderOptions}
              />
            </div>
          )}

          {/* Generate Button */}
          <button
            onClick={handleGenerate}
            disabled={loading || (editingImage ? !editInstructions.trim() : !prompt.trim())}
            className={`w-full flex items-center justify-center gap-2 disabled:opacity-50 disabled:cursor-not-allowed ${editingImage ? 'bg-purple-600 hover:bg-purple-700 text-white py-3 px-6 rounded-lg font-medium transition-colors' : 'btn-primary'
              }`}
          >
            {editingImage ? <Pencil className="w-5 h-5" /> : <Sparkles className="w-5 h-5" />}
            {loading ? (editingImage ? 'Editing...' : 'Generating...') : (editingImage ? 'Apply Edits' : 'Generate Images')}
          </button>

          {/* Job Progress */}
          {jobId && loading && (
            <JobProgress
              jobId={jobId}
              onComplete={handleJobComplete}
              onError={handleJobError}
            />
          )}
        </div>

        {/* Results */}
        <div>
          <h2 className="text-lg font-semibold text-white mb-4">Generated Images</h2>
          {generatedImages.length > 0 ? (
            <div className="space-y-4">
              {generatedImages.map((image) => (
                <div
                  key={image.id}
                  className="bg-forge-dark rounded-xl overflow-hidden border border-gray-800 group"
                >
                  <div className="relative w-full group">
                    <img
                      src={`/api/v1/assets/${image.id}/download`}
                      alt="Generated"
                      className="w-full h-auto object-contain bg-black/20"
                    />
                    <div className="absolute inset-0 bg-black/60 opacity-0 group-hover:opacity-100 transition-opacity flex flex-col items-center justify-center gap-2 p-4">
                      {supportsEditing && (
                        <button
                          onClick={() => handleStartEdit(image)}
                          className="bg-purple-600 hover:bg-purple-700 text-white py-1.5 px-3 rounded-lg font-medium text-sm transition-colors flex items-center gap-2 w-full justify-center"
                          title="Edit with Nano Banana"
                        >
                          <Pencil className="w-4 h-4" />
                          Edit
                        </button>
                      )}
                      <button
                        onClick={() => router.push(`/image/upscale?assetId=${image.id}`)}
                        className="bg-blue-600 hover:bg-blue-700 text-white py-1.5 px-3 rounded-lg font-medium text-sm transition-colors flex items-center gap-2 w-full justify-center"
                      >
                        <Maximize className="w-4 h-4" />
                        Upscale
                      </button>
                      <button
                        onClick={() => router.push(`/video/generate?assetId=${image.id}`)}
                        className="bg-green-600 hover:bg-green-700 text-white py-1.5 px-3 rounded-lg font-medium text-sm transition-colors flex items-center gap-2 w-full justify-center"
                      >
                        <Film className="w-4 h-4" />
                        Video
                      </button>
                      <button
                        onClick={() => handleDownload(image.id, image.original_filename)}
                        className="btn-primary py-1.5 px-3 text-sm flex items-center gap-2 w-full justify-center"
                      >
                        <Download className="w-4 h-4" />
                        Download
                      </button>
                    </div>
                  </div>
                </div>
              ))}
            </div>
          ) : (
            <div className="bg-forge-dark rounded-xl border border-gray-800 aspect-square flex items-center justify-center">
              <p className="text-gray-500">Generated images will appear here</p>
            </div>
          )}
        </div>
      </div>
    </div>
  );
}
