import gradio as gr
import torch
from transformers import CLIPProcessor, CLIPModel
import numpy as np
import kagglehub
from PIL import Image
import os
from pathlib import Path
import logging
import faiss
from tqdm import tqdm
import speech_recognition as sr
from gtts import gTTS
import tempfile

# Configure logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

class ImageSearchSystem:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        logger.info(f"Using device: {self.device}")
        
        # Load CLIP model
        self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
        self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16").to(self.device)
        
        # Initialize dataset
        self.image_paths = []
        self.index = None
        self.initialized = False
        
    def initialize_dataset(self) -> None:
        """Download and process dataset"""
        try:
            path = kagglehub.dataset_download("alessandrasala79/ai-vs-human-generated-dataset")
            image_folder = os.path.join(path, 'test_data_v2')
            
            self.image_paths = [
                f for f in Path(image_folder).glob("**/*") 
                if f.suffix.lower() in ['.jpg', '.jpeg', '.png']
            ]
            
            if not self.image_paths:
                raise ValueError(f"No images found in {image_folder}")
            
            logger.info(f"Found {len(self.image_paths)} images")
            
            self._create_image_index()
            self.initialized = True
            
        except Exception as e:
            logger.error(f"Dataset initialization failed: {str(e)}")
            raise

    def _create_image_index(self, batch_size: int = 32) -> None:
        """Create FAISS index"""
        try:
            all_features = []
            
            for i in tqdm(range(0, len(self.image_paths), batch_size), desc="Indexing images"):
                batch_paths = self.image_paths[i:i + batch_size]
                batch_images = [Image.open(img).convert("RGB") for img in batch_paths]

                if batch_images:
                    inputs = self.processor(images=batch_images, return_tensors="pt", padding=True)
                    inputs = {k: v.to(self.device) for k, v in inputs.items()}
                    
                    with torch.no_grad():
                        features = self.model.get_image_features(**inputs)
                        features = features / features.norm(dim=-1, keepdim=True)
                        all_features.append(features.cpu().numpy())
            
            all_features = np.concatenate(all_features, axis=0)
            self.index = faiss.IndexFlatIP(all_features.shape[1])
            self.index.add(all_features)
            
            logger.info("Image index created successfully")
            
        except Exception as e:
            logger.error(f"Failed to create image index: {str(e)}")
            raise

    def search(self, query: str, audio_path: str = None, k: int = 5):
        """Search for images using text or speech"""
        try:
            if not self.initialized:
                raise RuntimeError("System not initialized. Call initialize_dataset() first.")
            
            # Convert speech to text if audio input is provided
            if audio_path:
                recognizer = sr.Recognizer()
                with sr.AudioFile(audio_path) as source:
                    audio_data = recognizer.record(source)
                    try:
                        query = recognizer.recognize_google(audio_data)
                    except sr.UnknownValueError:
                        return [], "Could not understand the spoken query.", None

            # Process text query
            inputs = self.processor(text=[query], return_tensors="pt", padding=True)
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            
            with torch.no_grad():
                text_features = self.model.get_text_features(**inputs)
                text_features = text_features / text_features.norm(dim=-1, keepdim=True)
            
            # Search FAISS index
            scores, indices = self.index.search(text_features.cpu().numpy(), k)
            results = [Image.open(self.image_paths[idx]) for idx in indices[0]]

            # Generate Text-to-Speech
            tts = gTTS(f"Showing results for {query}")
            temp_audio = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
            tts.save(temp_audio.name)

            return results, query, temp_audio.name

        except Exception as e:
            logger.error(f"Search failed: {str(e)}")
            return [], "Error during search.", None

def create_demo_interface() -> gr.Interface:
    """Create Gradio interface with dark mode & speech support"""
    system = ImageSearchSystem()
    
    try:
        system.initialize_dataset()
    except Exception as e:
        logger.error(f"Failed to initialize system: {str(e)}")
        raise
    
    examples = [
        ["a beautiful landscape with mountains"],
        ["people working in an office"],
        ["a cute dog playing"],
        ["a modern city skyline at night"],
        ["a delicious-looking meal"]
    ]
    
    return gr.Interface(
        fn=system.search,
        inputs=[
            gr.Textbox(label="Enter your search query:", placeholder="Describe the image...", lines=2),
            gr.Audio(source="microphone", type="filepath", label="Speak Your Query (Optional)")
        ],
        outputs=[
            gr.Gallery(label="Search Results", show_label=True, columns=5, height="auto"),
            gr.Textbox(label="Spoken Query", interactive=False),
            gr.Audio(label="Results Spoken Out Loud")
        ],
        title="Multi-Modal Image Search",
        description="Use text or voice to search for images.",
        theme="dark",
        examples=examples,
        cache_examples=True,
        css=".gradio-container {background-color: #121212; color: #ffffff;}"
    )

if __name__ == "__main__":
    try:
        demo = create_demo_interface()
        demo.launch(share=True, enable_queue=True, max_threads=40)
    except Exception as e:
        logger.error(f"Failed to launch app: {str(e)}")
        raise