gradio_app.py

import os
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
import torch
from datetime import datetime
import gradio as gr
from typing import Dict, List, Union, Optional
import logging
import traceback
import asyncio
import httpx
import subprocess
import atexit

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Start the API server
def start_api_server():
    # Start uvicorn in a subprocess
    process = subprocess.Popen(["uvicorn", "script_search_api:app", "--reload"])
    return process

# Stop the API server
def stop_api_server(process):
    process.terminate()

# Register the exit handler
api_process = start_api_server()
atexit.register(stop_api_server, api_process)

class FlanT5Analyzer:
    """Fast and efficient analyzer using FLAN-T5"""
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = None
        self.tokenizer = None
        self.batch_size = 4
        self.trigger_categories = {
            "Violence": {
                "mapped_name": "Violence",
                "description": (
                    "Any act involving physical force or aggression intended to cause harm, injury, or death to a person, animal, or object. "
                    "Includes direct physical confrontations (e.g., fights, beatings, or assaults), implied violence (e.g., very graphical threats or descriptions of injuries), "
                    "or large-scale events like wars, riots, or violent protests."
                )
            },
            "Death": {
                "mapped_name": "Death References",
                "description": (
                    "Any mention, implication, or depiction of the loss of life, including direct deaths of characters, including mentions of deceased individuals, "
                    "or abstract references to mortality (e.g., 'facing the end' or 'gone forever'). This also covers depictions of funerals, mourning, "
                    "grieving, or any dialogue that centers around death, do not take metaphors into context that don't actually lead to death."
                )
            },
            "Substance_Use": {
                "mapped_name": "Substance Use",
                "description": (
                    "Any explicit reference to the consumption, misuse, or abuse of drugs, alcohol, or other intoxicating substances. "
                    "This includes scenes of drug use, drinking, smoking, discussions about heavy substance abuse or substance-related paraphernalia."
                )
            },
            "Gore": {
                "mapped_name": "Gore",
                "description": (
                    "Extremely detailed and graphic depictions of highly severe physical injuries, mutilation, or extreme bodily harm, often accompanied by descriptions of heavy blood, exposed organs, "
                    "or dismemberment. This includes war scenes with severe casualties, horror scenarios involving grotesque creatures, or medical procedures depicted with excessive detail."
                )
            },
            "Sexual_Content": {
                "mapped_name": "Sexual Content",
                "description": (
                    "Any depiction of sexual activity, intimacy, or sexual behavior, ranging from implied scenes to explicit descriptions. "
                    "This includes physical descriptions of characters in a sexual context, sexual dialogue, or references to sexual themes."
                )
            },
            "Sexual_Abuse": {
               "mapped_name": "Sexual Abuse",
               "description": (
                  "Any form of non-consensual sexual act, behavior, or interaction, involving coercion, manipulation, or physical force. "
                  "This includes incidents of sexual assault, exploitation, harassment, and any acts where an individual is subjected to sexual acts against their will."
               )
            },
            "Self_Harm": {
                "mapped_name": "Self-Harm",
                "description": (
                    "Any mention or depiction of behaviors where an individual intentionally causes harm to themselves. This includes cutting, burning, or other forms of physical injury, "
                    "as well as suicidal ideation, suicide attempts, or discussions of self-destructive thoughts and actions."
                )
            },
            "Mental_Health": {
                "mapped_name": "Mental Health Issues",
                "description": (
                    "Any reference to extreme mental health struggles, disorders, or psychological distress. This includes depictions of depression, anxiety, PTSD, bipolar disorder, "
                    "or other conditions. Also includes toxic traits such as Gaslighting or other psychological horrors"
                )
            }
        }
        logger.info(f"Initialized FLAN-T5 analyzer with device: {self.device}")

    async def load_model(self, progress=None) -> None:
        """Load the model and tokenizer with progress updates."""
        try:
            if progress:
                progress(0.1, "🍩 Loading tokenizer...")
            
            self.tokenizer = AutoTokenizer.from_pretrained(
                "google/flan-t5-base",
                use_fast=True
            )
            
            if progress:
                progress(0.3, "🍰 Loading model...")
            
            self.model = AutoModelForSeq2SeqLM.from_pretrained(
                "google/flan-t5-base",
                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
                device_map="auto"
            )
            
            if self.device == "cuda":
                self.model.eval()
                torch.cuda.empty_cache()
                
            if progress:
                progress(0.5, "🧁 Model loaded successfully")
                
        except Exception as e:
            logger.error(f"Error loading model: {str(e)}")
            raise

    def _chunk_text(self, text: str, chunk_size: int = 512, overlap: int = 30) -> List[str]:
        """Split text into overlapping chunks."""
        words = text.split()
        chunks = []
        for i in range(0, len(words), chunk_size - overlap):
            chunk = ' '.join(words[i:i + chunk_size])
            chunks.append(chunk)
        return chunks

    def _validate_response(self, response: str) -> str:
        """Validate and clean model response."""
        valid_responses = {"YES", "NO", "MAYBE"}
        response = response.strip().upper()
        first_word = response.split()[0] if response else "NO"
        return first_word if first_word in valid_responses else "NO"

    async def analyze_chunks_batch(
        self,
        chunks: List[str],
        progress: Optional[gr.Progress] = None,
        current_progress: float = 0,
        progress_step: float = 0
    ) -> Dict[str, float]:
        """Analyze multiple chunks in batches."""
        all_triggers = {}
        
        for category, info in self.trigger_categories.items():
            mapped_name = info["mapped_name"]
            description = info["description"]
            
            for i in range(0, len(chunks), self.batch_size):
                batch_chunks = chunks[i:i + self.batch_size]
                prompts = []
                
                for chunk in batch_chunks:
                    prompt = f"""
                    Task: Analyze if this text contains {mapped_name}.
                    Context: {description}
                    Text: "{chunk}"
                    
                    Rules for analysis:
                    1. Only answer YES if there is clear, direct evidence
                    2. Answer NO if the content is ambiguous or metaphorical
                    3. Consider the severity and context
                    
                    Answer with ONLY ONE word: YES, NO, or MAYBE
                    """
                    prompts.append(prompt)

                try:
                    inputs = self.tokenizer(
                        prompts,
                        return_tensors="pt",
                        padding=True,
                        truncation=True,
                        max_length=512
                    ).to(self.device)
                    
                    with torch.no_grad():
                        outputs = self.model.generate(
                            **inputs,
                            max_new_tokens=20,
                            temperature=0.2,
                            top_p=0.85,
                            num_beams=3,
                            early_stopping=True,
                            pad_token_id=self.tokenizer.eos_token_id,
                            do_sample=True
                        )
                    
                    responses = [
                        self.tokenizer.decode(output, skip_special_tokens=True)
                        for output in outputs
                    ]
                    
                    for response in responses:
                        validated_response = self._validate_response(response)
                        if validated_response == "YES":
                            all_triggers[mapped_name] = all_triggers.get(mapped_name, 0) + 1
                        elif validated_response == "MAYBE":
                            all_triggers[mapped_name] = all_triggers.get(mapped_name, 0) + 0.5
                
                except Exception as e:
                    logger.error(f"Error processing batch for {mapped_name}: {str(e)}")
                    continue
                
                if progress:
                    current_progress += progress_step
                    progress(min(current_progress, 0.9), f"🍭 Analyzing {mapped_name}...")
                    
        return all_triggers

    async def analyze_script(self, script: str, progress: Optional[gr.Progress] = None) -> List[str]:
        """Analyze the entire script."""
        if not self.model or not self.tokenizer:
            await self.load_model(progress)
        
        chunks = self._chunk_text(script)
        identified_triggers = await self.analyze_chunks_batch(
            chunks,
            progress,
            current_progress=0.5,
            progress_step=0.4 / (len(chunks) * len(self.trigger_categories))
        )
        
        if progress:
            progress(0.95, "🍫 Finalizing results...")

        final_triggers = []
        chunk_threshold = max(1, len(chunks) * 0.1)
        
        for mapped_name, count in identified_triggers.items():
            if count >= chunk_threshold:
                final_triggers.append(mapped_name)

        return final_triggers if final_triggers else ["None"]


class LlamaAnalyzer:
    """Detailed analyzer using Llama for thorough analysis"""
    def __init__(self):
        self.hf_token = os.getenv("HF_TOKEN")
        if not self.hf_token:
            raise ValueError("HF_TOKEN environment variable is not set!")
        
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = None
        self.tokenizer = None
        logger.info(f"Initialized Llama analyzer with device: {self.device}")

    async def load_model(self, progress=None) -> None:
        """Load the model and tokenizer with progress updates and detailed logging."""
        try:
            print("\n=== Starting Llama Model Loading ===")
            print(f"Time: {datetime.now()}")
            
            if progress:
                progress(0.1, "🍩 Loading Llama tokenizer...")
            
            print("Loading tokenizer...")
            self.tokenizer = AutoTokenizer.from_pretrained(
                "meta-llama/Llama-3.2-3B",
                use_fast=True
            )

            if progress:
                progress(0.3, "🍰 Loading Llama model...")
            
            print(f"Loading model on {self.device}...")
            self.model = AutoModelForCausalLM.from_pretrained(
                "meta-llama/Llama-3.2-3B",
                token=self.hf_token,
                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
                device_map="auto"
            )

            if progress:
                progress(0.5, "🧁 Llama model loaded successfully")
            
            print("Model and tokenizer loaded successfully")
            logger.info(f"Model loaded successfully on {self.device}")
        except Exception as e:
            logger.error(f"Error loading model: {str(e)}")
            print(f"\nERROR DURING MODEL LOADING: {str(e)}")
            print("Stack trace:")
            traceback.print_exc()
            raise

    def _chunk_text(self, text: str, chunk_size: int = 256, overlap: int = 15) -> List[str]:
        """Split text into overlapping chunks for processing."""
        chunks = []
        for i in range(0, len(text), chunk_size - overlap):
            chunk = text[i:i + chunk_size]
            chunks.append(chunk)
        print(f"Split text into {len(chunks)} chunks with {overlap} token overlap")
        return chunks

    async def analyze_chunk(
        self,
        chunk: str,
        trigger_categories: Dict,
        progress: Optional[gr.Progress] = None,
        current_progress: float = 0,
        progress_step: float = 0
    ) -> Dict[str, float]:
        """Analyze a single chunk of text for triggers with detailed logging."""
        chunk_triggers = {}
        print(f"\n--- Processing Chunk ---")
        print(f"Chunk text (preview): {chunk[:50]}...")
        
        for category, info in trigger_categories.items():
            mapped_name = info["mapped_name"]
            description = info["description"]

            print(f"\nAnalyzing for {mapped_name}...")
            prompt = f"""
            Check this text for any clear indication of {mapped_name} ({description}).
            only say yes if you are confident, make sure the text is not metaphorical.
            Respond concisely and only with: YES, NO, or MAYBE.
            Text: {chunk}
            Answer:
            """

            try:
                print("Sending prompt to model...")
                inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
                inputs = {k: v.to(self.device) for k, v in inputs.items()}

                with torch.no_grad():
                    print("Generating response...")
                    outputs = self.model.generate(
                        **inputs,
                        max_new_tokens=2,
                        do_sample=True,
                        temperature=0.3,
                        top_p=0.9,
                        pad_token_id=self.tokenizer.eos_token_id
                    )

                response_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True).strip().upper()
                first_word = response_text.split("\n")[-1].split()[0] if response_text else "NO"
                print(f"Model response for {mapped_name}: {first_word}")

                if first_word == "YES":
                    print(f"Detected {mapped_name} in this chunk!")
                    chunk_triggers[mapped_name] = chunk_triggers.get(mapped_name, 0) + 1
                elif first_word == "MAYBE":
                    print(f"Possible {mapped_name} detected, marking for further review.")
                    chunk_triggers[mapped_name] = chunk_triggers.get(mapped_name, 0) + 0.5
                else:
                    print(f"No {mapped_name} detected in this chunk.")

                if progress:
                    current_progress += progress_step
                    progress(min(current_progress, 0.9), f"🍭 Analyzing {mapped_name}...")

            except Exception as e:
                logger.error(f"Error analyzing chunk for {mapped_name}: {str(e)}")
                print(f"Error during analysis of {mapped_name}: {str(e)}")
                traceback.print_exc()

        return chunk_triggers

    async def analyze_script(self, script: str, progress: Optional[gr.Progress] = None) -> List[str]:
        """Analyze the entire script for triggers with progress updates and detailed logging."""
        print("\n=== Starting Script Analysis ===")
        print(f"Time: {datetime.now()}")

        if not self.model or not self.tokenizer:
            await self.load_model(progress)

        # Initialize trigger categories
        trigger_categories = {
            "Violence": {
                "mapped_name": "Violence",
                "description": (
                    "Any act of physical force meant to cause harm, injury, or death, including fights, threats, and large-scale violence like wars or riots."
                )
            },
            "Death": {
                "mapped_name": "Death References",
                "description": (
                    "Mentions or depictions of death, such as characters dying, references to deceased people, funerals, or mourning."
                )
            },
            "Substance Use": {
                "mapped_name": "Substance Use",
                "description": (
                    "Any reference to using or abusing drugs, alcohol, or other substances, including scenes of drinking, smoking, or drug use."
                )
            },
            "Gore": {
                "mapped_name": "Gore",
                "description": (
                    "Graphic depictions of severe injuries or mutilation, often with detailed blood, exposed organs, or dismemberment."
                )
            },
            "Vomit": {
                "mapped_name": "Vomit",
                "description": (
                    "Any explicit reference to vomiting or related actions. This includes only very specific mentions of nausea or the act of vomiting, with more focus on the direct description, only flag this if you absolutely believe it's present."
                )
            },
            "Sexual Content": {
                "mapped_name": "Sexual Content",
                "description": (
                    "Depictions or mentions of sexual activity, intimacy, or behavior, including sexual themes like harassment or innuendo."
                )
            },
            "Sexual Abuse": {
                "mapped_name": "Sexual Abuse",
                "description": (
                    "Explicit non-consensual sexual acts, including assault, molestation, or harassment, and the emotional or legal consequences of such abuse. A stronger focus on detailed depictions or direct references to coercion or violence."
                )
            },
            "Self-Harm": {
                "mapped_name": "Self-Harm",
                "description": (
                    "Depictions or mentions of intentional self-injury, including acts like cutting, burning, or other self-destructive behavior. Emphasis on more graphic or repeated actions, not implied or casual references."
                )
            },
            "Gun Use": {
                "mapped_name": "Gun Use",
                "description": (
                    "Explicit mentions of firearms in use, including threatening actions or accidents involving guns. Only triggers when the gun use is shown in a clear, violent context."
                )
            },
            "Animal Cruelty": {
                "mapped_name": "Animal Cruelty",
                "description": (
                    "Direct or explicit harm, abuse, or neglect of animals, including physical abuse or suffering, and actions performed for human entertainment or experimentation. Triggers only in clear, violent depictions."
                )
            },
            "Mental Health Issues": {
                "mapped_name": "Mental Health Issues",
                "description": (
                    "References to psychological struggles, such as depression, anxiety, or PTSD, including therapy or coping mechanisms."
                )
            }
        }

        chunks = self._chunk_text(script)
        identified_triggers = {}
        progress_step = 0.4 / (len(chunks) * len(trigger_categories))
        current_progress = 0.5  # Starting after model loading

        for chunk_idx, chunk in enumerate(chunks, 1):
            chunk_triggers = await self.analyze_chunk(
                chunk,
                trigger_categories,
                progress,
                current_progress,
                progress_step
            )
            
            for trigger, count in chunk_triggers.items():
                identified_triggers[trigger] = identified_triggers.get(trigger, 0) + count

        if progress:
            progress(0.95, "🍫 Finalizing detailed results...")

        print("\n=== Analysis Complete ===")
        print("Final Results:")
        final_triggers = []

        for mapped_name, count in identified_triggers.items():
            if count > 0.5:
                final_triggers.append(mapped_name)
            print(f"- {mapped_name}: found in {count} chunks")

        if not final_triggers:
            print("No triggers detected")
            final_triggers = ["None"]

        return final_triggers


async def analyze_content_flant5(
    script: str,
    progress: Optional[gr.Progress] = None
) -> Dict[str, Union[List[str], str]]:
    """Main analysis function using FLAN-T5."""
    logger.info("Starting FLAN-T5 content analysis")
    
    analyzer = FlanT5Analyzer()
    
    try:
        triggers = await analyzer.analyze_script(script, progress)
        
        if progress:
            progress(1.0, "🎉 Analysis complete!")

        result = {
            "detected_triggers": triggers,
            "confidence": "High - Content detected" if triggers != ["None"] else "High - No concerning content detected",
            "model": "google/flan-t5-base",
            "analysis_timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }

        logger.info(f"Analysis complete: {result}")
        return result

    except Exception as e:
        logger.error(f"Analysis error: {str(e)}")
        return {
            "detected_triggers": ["Error occurred during analysis"],
            "confidence": "Error",
            "model": "google/flan-t5-base",
            "analysis_timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "error": str(e)
        }


async def analyze_content_llama(
    script: str,
    progress: Optional[gr.Progress] = None
) -> Dict[str, Union[List[str], str]]:
    """Main analysis function using Llama for detailed analysis."""
    print("\n=== Starting Llama Content Analysis ===")
    print(f"Time: {datetime.now()}")
    
    analyzer = LlamaAnalyzer()
    
    try:
        triggers = await analyzer.analyze_script(script, progress)
        
        if progress:
            progress(1.0, "🎉 Detailed analysis complete!")

        result = {
            "detected_triggers": triggers,
            "confidence": "High - Content detected" if triggers != ["None"] else "High - No concerning content detected",
            "model": "Llama-3.2-3B",
            "analysis_timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }

        print("\nFinal Result Dictionary:", result)
        return result

    except Exception as e:
        logger.error(f"Analysis error: {str(e)}")
        print(f"\nERROR OCCURRED: {str(e)}")
        print("Stack trace:")
        traceback.print_exc()
        return {
            "detected_triggers": ["Error occurred during analysis"],
            "confidence": "Error",
            "model": "Llama-3.2-3B",
            "analysis_timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "error": str(e)
        }


async def search_movie_script(movie_name: str) -> Dict:
    """Search for and analyze a movie script using the API service."""
    async with httpx.AsyncClient() as client:
        try:
            # Start search
            response = await client.post(
                "http://localhost:8000/search",
                json={"movie_name": movie_name}
            )
            response.raise_for_status()
            task_data = response.json()
            task_id = task_data["task_id"]
            
            # Poll for results
            while True:
                status_response = await client.get(f"http://localhost:8000/progress/{task_id}")
                status_response.raise_for_status()
                status_data = status_response.json()
                
                if status_data["is_complete"]:
                    if status_data["error"]:
                        return {"error": status_data["error"]}
                    return status_data["result"]
                
                await asyncio.sleep(1)
                
        except Exception as e:
            return {"error": f"Error during movie search: {str(e)}"}


# CSS animations for movie search tab
movie_search_html = """
    <div style="text-align: center; padding: 15px;">
        <h3>🎬 Movie Script Analysis</h3>
        <p style="color: #ffffff;">Search and analyze movie scripts for content warnings</p>
    </div>
"""

# Custom CSS for Monochrome theme with subtle animations
custom_css = """
/* Monochrome Black Theme */
body, .gradio-container {
    background: #000000 !important;
    color: #ffffff !important;
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}

/* Subtle animated background */
.gradio-container::before {
    content: '';
    position: fixed;
    top: 0;
    left: 0;
    width: 100%;
    height: 100%;
    background: 
        radial-gradient(circle at 20% 20%, rgba(255, 255, 255, 0.02) 0%, transparent 30%),
        radial-gradient(circle at 80% 80%, rgba(255, 255, 255, 0.02) 0%, transparent 30%),
        radial-gradient(circle at 40% 60%, rgba(255, 255, 255, 0.01) 0%, transparent 40%);
    animation: subtle-float 30s ease-in-out infinite;
    pointer-events: none;
    z-index: -1;
}

@keyframes subtle-float {
    0%, 100% { transform: translateY(0px) rotate(0deg); }
    50% { transform: translateY(-10px) rotate(1deg); }
}

/* Loading animation */
@keyframes spin {
    0% { transform: rotate(0deg); }
    100% { transform: rotate(360deg); }
}

.loading-spinner {
    display: inline-block;
    width: 40px;
    height: 40px;
    border: 4px solid rgba(255, 255, 255, 0.3);
    border-radius: 50%;
    border-top-color: #ffffff;
    animation: spin 1s ease-in-out infinite;
}

/* Progress bar styling */
.progress-bar {
    background: linear-gradient(90deg, #333333, #ffffff, #333333) !important;
    border-radius: 20px !important;
    animation: progress-pulse 2s ease-in-out infinite alternate;
}

@keyframes progress-pulse {
    from { opacity: 0.7; }
    to { opacity: 1; }
}

/* Input and output styling */
.gr-textbox, .gr-json {
    background: rgba(255, 255, 255, 0.05) !important;
    border: 2px solid rgba(255, 255, 255, 0.2) !important;
    border-radius: 8px !important;
    color: #ffffff !important;
    transition: all 0.3s ease !important;
}

.gr-textbox:focus, .gr-json:focus {
    border-color: rgba(255, 255, 255, 0.4) !important;
    box-shadow: 0 0 10px rgba(255, 255, 255, 0.1) !important;
    background: rgba(255, 255, 255, 0.08) !important;
}

/* Button styling */
.gr-button {
    background: linear-gradient(45deg, #333333, #555555) !important;
    border: 2px solid rgba(255, 255, 255, 0.2) !important;
    border-radius: 8px !important;
    color: #ffffff !important;
    font-weight: bold !important;
    padding: 12px 24px !important;
    transition: all 0.3s ease !important;
    position: relative !important;
    overflow: hidden !important;
}

.gr-button::before {
    content: '';
    position: absolute;
    top: 50%;
    left: 50%;
    width: 0;
    height: 0;
    background: rgba(255, 255, 255, 0.1);
    border-radius: 50%;
    transform: translate(-50%, -50%);
    transition: width 0.6s, height 0.6s;
}

.gr-button:hover::before {
    width: 300px;
    height: 300px;
}

.gr-button:hover {
    background: linear-gradient(45deg, #555555, #777777) !important;
    border-color: rgba(255, 255, 255, 0.4) !important;
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 16px rgba(0, 0, 0, 0.3) !important;
}

/* Tab styling */
.gr-tab {
    background: rgba(255, 255, 255, 0.05) !important;
    border: 2px solid rgba(255, 255, 255, 0.2) !important;
    border-radius: 8px 8px 0 0 !important;
    color: #ffffff !important;
    transition: all 0.3s ease !important;
}

.gr-tab:hover, .gr-tab.selected {
    background: rgba(255, 255, 255, 0.1) !important;
    border-color: rgba(255, 255, 255, 0.4) !important;
    transform: translateY(-2px) !important;
}

/* Title styling */
h1, h2, h3 {
    color: #ffffff !important;
    text-align: center;
    margin: 20px 0;
    text-shadow: 0 0 10px rgba(255, 255, 255, 0.3);
    animation: title-fade 4s ease-in-out infinite alternate;
}

@keyframes title-fade {
    from { opacity: 0.9; }
    to { opacity: 1; }
}

/* Emoji animations */
.sweet-emoji {
    display: inline-block;
    animation: gentle-bounce 3s ease-in-out infinite;
}

@keyframes gentle-bounce {
    0%, 100% { transform: translateY(0); }
    50% { transform: translateY(-5px); }
}

/* Container styling */
.gr-group {
    background: rgba(255, 255, 255, 0.03) !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
    border-radius: 8px !important;
    box-shadow: 0 4px 16px rgba(0, 0, 0, 0.2) !important;
}

/* JSON output styling */
.gr-json pre {
    background: rgba(255, 255, 255, 0.05) !important;
    border: 1px solid rgba(255, 255, 255, 0.2) !important;
    border-radius: 8px !important;
    color: #ffffff !important;
    font-family: 'Fira Code', monospace !important;
}

/* Progress text styling */
.progress-text {
    color: #ffffff !important;
    font-weight: bold !important;
    text-shadow: 0 0 5px rgba(255, 255, 255, 0.3) !important;
}

/* Label styling */
label {
    color: #ffffff !important;
}

/* Paragraph styling */
p {
    color: #ffffff !important;
}

/* Ensure all text is white */
* {
    color: #ffffff !important;
}

/* Special styling for colored text elements */
.gr-textbox textarea, .gr-textbox input {
    color: #ffffff !important;
    background: transparent !important;
}

/* Placeholder text styling */
.gr-textbox textarea::placeholder, .gr-textbox input::placeholder {
    color: rgba(255, 255, 255, 0.5) !important;
}
"""

if __name__ == "__main__":
    # Create the Gradio interface with tabs
    with gr.Blocks(
        title="🍩 TREAT-CHOCOSYRUP 🧁 Content Analysis", 
        css=custom_css,
        theme=gr.themes.Base()
    ) as app:
        gr.HTML("""
            <div style="text-align: center; padding: 20px;">
                <h1 style="font-size: 3em; margin-bottom: 10px;">
                    🍩 TREAT-CHOCOSYRUP 🧁
                </h1>
                <h2 style="font-size: 1.5em; margin-top: 0;">
                    <span class="sweet-emoji">🍭</span> Content Analysis & Trigger Detection <span class="sweet-emoji">🍰</span>
                </h2>
            </div>
        """)
        
        with gr.Tabs():
            # Default FLAN-T5 Tab
            with gr.Tab("🍰 Quick Analysis (FLAN-T5)", elem_id="flant5-tab"):
                gr.HTML("""
                    <div style="text-align: center; padding: 15px;">
                        <h3>🚀 Fast & Efficient Analysis</h3>
                        <p style="color: #ff9ff3;">Perfect for quick content screening with high accuracy</p>
                    </div>
                """)

                # Input and analyze button for FLAN-T5
                with gr.Row():
                    text_input_flant5 = gr.Textbox(
                        lines=8,
                        label="Input Text",
                        placeholder="Enter your text here for analysis...",
                        elem_id="input-text-flant5"
                    )

                with gr.Row():
                    analyze_button_flant5 = gr.Button("🔍 Analyze Content", variant="primary")
                    clear_button_flant5 = gr.Button("🧹 Clear", variant="secondary")

                # Output area for FLAN-T5
                output_json_flant5 = gr.JSON(label="Analysis Results")

                # Button click events for FLAN-T5
                analyze_button_flant5.click(
                    fn=analyze_content_flant5,
                    inputs=[text_input_flant5],
                    outputs=[output_json_flant5]
                )
                clear_button_flant5.click(
                    fn=lambda: ("", None),
                    inputs=[],
                    outputs=[text_input_flant5, output_json_flant5]
                )

            # Detailed Llama Analysis Tab
            with gr.Tab("🍬 Detailed Analysis (Llama)", elem_id="llama-tab"):
                gr.HTML("""
                    <div style="text-align: center; padding: 15px;">
                        <h3>🎯 Deep & Thorough Analysis</h3>
                        <p style="color: #4ecdc4;">Advanced analysis for comprehensive content evaluation</p>
                    </div>
                """)

                # Input and analyze button for Llama
                with gr.Row():
                    text_input_llama = gr.Textbox(
                        lines=8,
                        label="Input Text",
                        placeholder="Enter your text here for detailed analysis...",
                        elem_id="input-text-llama"
                    )

                with gr.Row():
                    analyze_button_llama = gr.Button("🔍 Analyze Content (Detailed)", variant="primary")
                    clear_button_llama = gr.Button("🧹 Clear", variant="secondary")

                # Output area for Llama
                output_json_llama = gr.JSON(label="Detailed Analysis Results")

                # Button click events for Llama
                analyze_button_llama.click(
                    fn=analyze_content_llama,
                    inputs=[text_input_llama],
                    outputs=[output_json_llama]
                )
                clear_button_llama.click(
                    fn=lambda: ("", None),
                    inputs=[],
                    outputs=[text_input_llama, output_json_llama]
                )

            # Movie Search and Analysis Tab
            with gr.Tab("🍿 Agentic Movie Script Analysis", elem_id="movie-tab"):
                gr.HTML(movie_search_html)

                # Input for movie name
                with gr.Row():
                    movie_name_input = gr.Textbox(
                        lines=1,
                        label="Movie Name",
                        placeholder="Enter the movie name to search and analyze...",
                        elem_id="movie-name-input"
                    )

                # Analyze button for movie search
                with gr.Row():
                    search_button = gr.Button("🔍 Search & Analyze Movie", variant="primary")
                    clear_movie_button = gr.Button("🧹 Clear", variant="secondary")

                # Output area for movie analysis results
                output_movie_json = gr.JSON(label="Movie Analysis Results")

                # Button click events for movie search
                search_button.click(
                    fn=search_movie_script,
                    inputs=[movie_name_input],
                    outputs=[output_movie_json]
                )
                clear_movie_button.click(
                    fn=lambda: ("", None),
                    inputs=[],
                    outputs=[movie_name_input, output_movie_json]
                )

    app.queue()
    app.launch()