import { useState, useEffect } from 'react';
import { z } from "zod";
import { Users } from 'lucide-react';
import { toast } from 'sonner';
import { useLocation, useNavigate } from 'react-router-dom';
import { Progress } from "@/components/ui/progress";

import AIRecruiterForm, { formSchema } from './ai-recruiter/AIRecruiterForm';
import PersonaReviewList from './ai-recruiter/PersonaReviewList';
import { generateSyntheticPersonas } from '@/utils/personaGenerator';
import { usePersonaStorage, GENERATED_PERSONAS_KEY } from '@/hooks/usePersonaStorage';
import { Persona } from "@/types/persona";

interface AIRecruiterProps {
  targetFolderId?: string | null;
  targetFolderName?: string | null;
}

export default function AIRecruiter({ targetFolderId, targetFolderName }: AIRecruiterProps) {
  const location = useLocation();
  const navigate = useNavigate();
  const { loadPersonas, savePersonas } = usePersonaStorage();
  
  const [isGenerating, setIsGenerating] = useState(false);
  const [generatedPersonas, setGeneratedPersonas] = useState<Persona[]>([]);
  const [selectedPersonas, setSelectedPersonas] = useState<string[]>([]);
  const [showReview, setShowReview] = useState(false);
  const [generationProgress, setGenerationProgress] = useState(0);
  
  // Check URL params for state restoration
  useEffect(() => {
    const searchParams = new URLSearchParams(location.search);
    const mode = searchParams.get('mode');
    const tab = searchParams.get('tab');
    const step = searchParams.get('step');
    
    // If we're returning to the review step, restore the state
    if (mode === 'create' && tab === 'ai' && step === 'review') {
      const personas = loadPersonas();
      if (personas.length > 0) {
        setGeneratedPersonas(personas);
        setSelectedPersonas(personas.map(p => p.id));
        setShowReview(true);
      }
    }
  }, [location, loadPersonas]);
  
  async function onSubmit(values: z.infer<typeof formSchema>) {
    try {
      setIsGenerating(true);
      setGenerationProgress(0);
      
      // Validate count before proceeding
      const count = parseInt(values.personaCount);
      if (isNaN(count) || count < 1 || count > 10) {
        toast.error("Invalid number of personas", {
          description: "Please enter a number between 1 and 10"
        });
        setIsGenerating(false);
        return;
      }
      
      // Start progress animation
      setGenerationProgress(5); // Initial progress
      
      // Simulate progress while waiting for personas to generate
      const progressInterval = setInterval(() => {
        setGenerationProgress(prev => {
          // Increase gradually but never reach 100% until actually complete
          if (prev < 90) {
            return prev + Math.random() * 5;
          }
          return prev;
        });
      }, 500);
      
      // Adjust the expected time based on the count
      const estimatedTime = count <= 2 ? "30-60 seconds" : 
                           count <= 4 ? "1-2 minutes" : 
                           count <= 6 ? "2-3 minutes" :
                                       "3-5 minutes";
                                       
      // Warn about potential timeouts for larger counts
      if (count > 4) {
        toast.info("Generation may take longer", {
          description: `Generating ${count} personas at once may result in some timeouts. If this happens, the successfully created personas will still be saved.`,
          duration: 8000
        });
      }
      
      toast.info("Generating AI personas in parallel", {
        description: `Creating ${count} synthetic personas based on your brief. This may take ${estimatedTime}. Please be patient.`,
        duration: 10000
      });
      
      // Show folder info in toast if available
      if (targetFolderId && targetFolderName) {
        console.log(`Target folder for new personas: ID=${targetFolderId}, Name=${targetFolderName}`);
        toast.info(`Creating personas in "${targetFolderName}" folder`, {
          duration: 3000
        });
      } else {
        console.log("No target folder specified for new personas");
      }
      
      // Log which model is being used for generation
      console.log(`🤖 Starting persona generation with model: ${values.llm_model || 'gemini-2.5-pro'}`);
      
      const response = await generateSyntheticPersonas(
        values.audienceBrief, 
        values.researchObjective,
        count,
        values.dataFile,
        targetFolderId,
        values.llm_model
      );
      
      // Extract personas from the response
      const personas = response.personas || response;
      
      // Clear the progress interval
      clearInterval(progressInterval);
      // Set progress to 100% when done
      setGenerationProgress(100);
      
      // Check for partial success (some personas generated, some failed)
      if (personas && personas.length > 0) {
        // Log successful generation with model info
        console.log(`✅ Successfully generated ${personas.length} personas using model: ${values.llm_model || 'gemini-2.5-pro'}`);
        
        // Check if we got a response with partial success info
        if (response.partial_success || (response.errors && response.errors.length > 0)) {
          // Some personas succeeded but others failed
          toast.success("Some personas generated successfully", {
            description: `${personas.length} synthetic personas were created using ${values.llm_model || 'Gemini 2.5 Pro'}. ${response.errors?.length || 0} failed due to timeout or other errors.`,
            duration: 8000
          });
          
          // Show details about the failures
          if (response.errors && response.errors.length > 0) {
            setTimeout(() => {
              toast.error("Some personas failed to generate", {
                description: `${response.errors.length} personas timed out. The server took too long to generate them. The successfully generated personas have been saved${targetFolderId ? ` in the selected folder` : ''}.`,
                duration: 10000
              });
            }, 1000);
          }
        } else {
          // All personas succeeded
          toast.success("Personas generated and saved successfully", {
            description: `${personas.length} synthetic personas have been created using ${values.llm_model || 'Gemini 2.5 Pro'} and saved ${targetFolderId ? `to the "${targetFolderName}" folder` : 'to the database'}.`
          });
        }
        
        // Navigate directly back to synthetic users list
        navigate('/synthetic-users?mode=view');
      } else {
        throw new Error("No personas were generated");
      }
    } catch (error) {
      console.error(`❌ Error generating personas using model: ${values.llm_model || 'gemini-2.5-pro'}:`, error);
      
      let errorMessage = "Please try again or adjust your parameters";
      let errorTitle = "Failed to generate personas";
      
      // Check for specific error types
      if (error.code === "ECONNABORTED" || error.message?.includes("timeout") || error.response?.status === 504) {
        errorTitle = "Generation timeout";
        errorMessage = "AI persona generation timed out. This often happens when generating multiple complex personas. Try generating fewer personas (2-3) or try again later.";
      } else if (error.response?.status === 500) {
        errorTitle = "Server error";
        
        // Try to extract the error message from the response if available
        if (error.response?.data?.message) {
          errorMessage = error.response.data.message;
        } else if (error.response?.data?.error) {
          errorMessage = error.response.data.error;
        } else {
          errorMessage = "The server encountered an error processing your request. Please try again later.";
        }
      } else if (error.response?.status === 401) {
        errorTitle = "Authentication required";
        errorMessage = "Please log in to generate personas.";
      } else if (error.message?.includes("504 Deadline Exceeded")) {
        errorTitle = "Generation timeout";
        errorMessage = "The AI model took too long to generate personas. Try generating fewer personas or simplify your brief.";
      } else if (error instanceof Error) {
        errorMessage = error.message;
      }
      
      toast.error(errorTitle, {
        description: errorMessage,
        duration: 6000 // Show for longer to ensure user sees it
      });
    } finally {
      // Wait a moment to show 100% completion before resetting
      setTimeout(() => {
        setIsGenerating(false);
        setGenerationProgress(0);
      }, 500);
    }
  }

  const handlePersonaSelection = (id: string) => {
    setSelectedPersonas(prev => 
      prev.includes(id) 
        ? prev.filter(pId => pId !== id)
        : [...prev, id]
    );
  };

  const processRefinementInstructions = (personas: Persona[], instructions: string) => {
    // This would connect to a real LLM in production
    // For demo purposes, we'll simulate some simple transformations
    
    const lowerInstructions = instructions.toLowerCase();
    
    return personas.map(p => {
      const persona = { ...p };
      
      // Handle age modifications
      if (lowerInstructions.includes('younger')) {
        const currentAge = parseInt(persona.age);
        persona.age = (currentAge - 5).toString(); // Make 5 years younger
      } else if (lowerInstructions.includes('older')) {
        const currentAge = parseInt(persona.age);
        persona.age = (currentAge + 5).toString(); // Make 5 years older
      }
      
      // Handle location diversity
      if (lowerInstructions.includes('different locations')) {
        // In a real implementation, this would be handled by the LLM
        // For demo, we'll just append a modifier to make it clear the change happened
        persona.location = `${persona.location} (Diversified)`;
      }
      
      // Handle personality traits
      if (lowerInstructions.includes('more extroverted')) {
        persona.personality = `Extroverted, ${persona.personality.toLowerCase()}`;
      } else if (lowerInstructions.includes('more introverted')) {
        persona.personality = `Introverted, ${persona.personality.toLowerCase()}`;
      }
      
      // Add some randomized modifications based on the instructions
      if (lowerInstructions.includes('diverse')) {
        // Randomize some traits
        const randomTraits = ['tech-savvy', 'traditional', 'innovative', 'conservative', 'creative'];
        const selectedTrait = randomTraits[Math.floor(Math.random() * randomTraits.length)];
        persona.personality = `${selectedTrait}, ${persona.personality}`;
      }
      
      return persona;
    });
  };

  const handleRefinePersonas = (refinementPrompt: string) => {
    if (!refinementPrompt.trim()) {
      toast.error("Please provide refinement instructions");
      return;
    }
    
    setIsGenerating(true);
    
    // In a production environment, this would call an actual LLM API
    setTimeout(() => {
      try {
        const selectedPersonaObjects = generatedPersonas.filter(p => 
          selectedPersonas.includes(p.id)
        );
        
        const refinedSelected = processRefinementInstructions(
          selectedPersonaObjects, 
          refinementPrompt
        );
        
        // Replace only the selected personas with their refined versions
        const updatedPersonas = generatedPersonas.map(p => {
          const refined = refinedSelected.find(r => r.id === p.id);
          return refined || p;
        });
        
        setGeneratedPersonas(updatedPersonas);
        setIsGenerating(false);
        savePersonas(updatedPersonas);
        
        toast.success("Personas refined based on your instructions", {
          description: "Review the updated profiles",
        });
      } catch (error) {
        console.error("Error refining personas:", error);
        toast.error("Failed to refine personas", {
          description: "Please try different instructions",
        });
        setIsGenerating(false);
      }
    }, 1500);
  };

  const handleApprovePersonas = () => {
    const approved = generatedPersonas.filter(p => selectedPersonas.includes(p.id));
    
    toast.success(`${approved.length} personas approved`, {
      description: "Added to your synthetic persona library",
    });
    
    savePersonas(approved);
    
    // Redirect to the persona library view instead of resetting
    navigate('/synthetic-users?mode=view');
  };

  return (
    <div className="glass-panel rounded-xl p-6">
      <div className="flex items-center gap-2 mb-6">
        <Users className="h-5 w-5 text-primary" />
        <h2 className="font-sf text-xl font-semibold">AI Persona Recruiter</h2>
      </div>
      
      {isGenerating && (
        <div className="mb-6">
          <div className="flex justify-between items-center mb-2">
            <span className="text-sm font-medium">Generating personas in parallel...</span>
            <span className="text-sm text-muted-foreground">{Math.round(generationProgress)}%</span>
          </div>
          <Progress value={generationProgress} className="h-2" />
        </div>
      )}
      
      {!showReview ? (
        <AIRecruiterForm 
          onSubmit={onSubmit}
          isGenerating={isGenerating}
        />
      ) : (
        <PersonaReviewList
          generatedPersonas={generatedPersonas}
          selectedPersonas={selectedPersonas}
          isGenerating={isGenerating}
          onPersonaSelection={handlePersonaSelection}
          onRefinePersonas={handleRefinePersonas}
          onApprovePersonas={handleApprovePersonas}
          onBackToGenerator={() => setShowReview(false)}
        />
      )}
    </div>
  );
}