app.py

# app.py

import gradio as gr
import torch
import logging
import re
import random
from gradio.themes import Soft
import json
import spaces
import gc
from typing import List, Dict, Tuple, Optional, Any
from dataclasses import dataclass
from transformers import (
    pipeline,
    MarianMTModel,
    MarianTokenizer,
    AutoTokenizer,
    AutoModelForSeq2SeqLM
)
from diffusers import StableDiffusionPipeline
import numpy as np
from PIL import Image

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

@dataclass
class StoryTemplate:
    """Story template data structure"""
    name: str
    template_en: str
    template_ar: str
    parameters: List[str]

class MultilingualStoryGenerator:
    """Robust story generation in both English and Arabic with content safety"""
    
    def __init__(self):
        try:
            # Initialize English generator
            self.en_generator = pipeline(
                "text-generation",
                model="EleutherAI/gpt-neo-1.3B",
                device="cpu"
            )
            logger.info("Initialized English story generator")
            
            # Initialize dedicated English sentiment analyzer
            self.en_sentiment = pipeline(
                "sentiment-analysis",
                model="distilbert-base-uncased-finetuned-sst-2-english",
                device="cpu"
            )
            logger.info("Initialized English sentiment analyzer")
            
            # Initialize Arabic MT5 model - better for controlled generation
            try:
                self.ar_tokenizer = AutoTokenizer.from_pretrained("google/mt5-small")
                self.ar_model = AutoModelForSeq2SeqLM.from_pretrained("google/mt5-small")
                logger.info("Initialized MT5 for Arabic generation")
                self.ar_model_available = True
            except Exception as e:
                logger.error(f"Failed to load Arabic MT5 model: {str(e)}")
                self.ar_model_available = False
                
            # Initialize Arabic sentiment analyzer if possible
            try:
                self.ar_sentiment = pipeline(
                    "sentiment-analysis",
                    model="CAMeL-Lab/bert-base-arabic-sentiment",
                    device="cpu"
                )
                logger.info("Initialized Arabic sentiment analyzer")
                self.ar_sentiment_available = True
            except Exception as e:
                logger.error(f"Failed to load Arabic sentiment: {str(e)}")
                self.ar_sentiment_available = False
                
            # Initialize translator models
            self.translator_ar_to_en = pipeline(
                "translation",
                model="Helsinki-NLP/opus-mt-ar-en",
                device="cpu"
            )
            logger.info("Initialized Arabic to English translator")
            
            # Attempt to load English to Arabic translator
            try:
                self.translator_en_to_ar = pipeline(
                    "translation",
                    model="Helsinki-NLP/opus-mt-en-ar",
                    device="cpu"
                )
                logger.info("Initialized English to Arabic translator")
                self.en_to_ar_available = True
            except Exception as e:
                logger.error(f"Failed to load EN->AR translator: {str(e)}")
                self.en_to_ar_available = False
                
            # Content filtering regex patterns
            self.banned_patterns = [
                r'sex', r'fuck', r'porn', r'explicit',
                r'malay', r'panama', r'solar', r'queent',
                # Arabic inappropriate terms
                r'جنس', r'فاحش', r'إباحي'
            ]
            
            # Pre-written high-quality templates
            self.story_templates = self._load_templates()
            
            self.initialized = True
        except Exception as e:
            logger.error(f"Failed to initialize MultilingualStoryGenerator: {str(e)}")
            self.initialized = False
    
    def _load_templates(self) -> Dict[str, Dict[str, List[str]]]:
        """Load high-quality story templates in both languages"""
        return {
            "english": {
                "adventure": [
                    "In the ancient ruins of {location}, a brave {hero} discovered a map leading to {object}. The journey wouldn't be easy, as dangerous traps and mythical creatures guarded the path. Armed with courage and determination, {hero} ventured into the unknown. After overcoming numerous challenges, including a treacherous river crossing and a riddle-speaking sphinx, {hero} finally reached the chamber where {object} was kept. This treasure wasn't just valuable - it held the power to change the world.",
                    
                    "The small village of {location} had a legend about {object}, a mysterious artifact said to grant its finder extraordinary powers. {hero}, a curious and adventurous soul, had always been fascinated by this tale. When strange events began occurring in the village, {hero} realized the time had come to seek {object}. The journey through dark forests and misty mountains tested every skill {hero} possessed. But with each obstacle overcome, {hero} grew stronger and wiser, eventually discovering that the true power of {object} was not what anyone had expected.",
                    
                    "{hero} had spent years studying ancient texts about {object}, hidden somewhere in {location}. Scholars had dismissed these stories as myths, but {hero} knew better. When troubling signs appeared in the sky, {hero} recognized them from the texts - it was time to find {object}. The expedition into {location} revealed dangers and wonders beyond imagination. {hero} faced not only physical challenges but moral dilemmas that questioned everything {hero} believed. In the final chamber, {hero} discovered that {object} was actually a gateway to knowledge that would forever change humanity's understanding of the universe."
                ],
                "friendship": [
                    "In a quiet neighborhood, a small cat named {character1} lived alone, scavenging for food and shelter. One rainy day, {character1} discovered an abandoned robot called {character2} in an alley, damaged and powered down. Curious, {character1} stayed nearby until a kind engineer reactivated {character2}. As {character2} came back to life, it was confused and disoriented. {character1}, despite being initially cautious, started visiting {character2} daily. Slowly, an unusual friendship formed. {character2} would protect {character1} from the neighborhood dogs, while {character1} would guide {character2} through the city streets. They taught each other about their very different perspectives on life, proving that friendship can transcend even the greatest differences.",
                    
                    "{character1} was the most advanced robot created by Future Technologies, designed to learn and evolve. However, {character1} felt something was missing in its programmed existence. One day, a stray cat named {character2} wandered into the laboratory. Instead of shooing it away, {character1} offered it food. Day after day, {character2} returned, and {character1} began to experience emotions not in its programming. When the company decided to reset {character1}'s learning algorithms, {character2} somehow sensed the danger and created a distraction that allowed {character1} to escape. Together, they embarked on an adventure that would redefine what it means to be alive and to care for another being.",
                    
                    "After a technical malfunction, {character1}, an experimental domestic robot, escaped from the research facility and hid in an abandoned building. There, it encountered {character2}, a street-smart cat who initially hissed and clawed at the strange metal creature. But when winter came and temperatures dropped, {character1} used its internal heating to keep {character2} warm. Gradually, {character2} began to trust the gentle robot. When researchers finally tracked down {character1}, they were amazed to find it had developed emotional responses through its friendship with {character2}. Instead of reclaiming the robot, they decided to study this unique bond, learning valuable lessons about connection and compassion that would transform robotics forever."
                ],
                "fantasy": [
                    "In the enchanted kingdom of {place}, young {protagonist} lived an ordinary life until the day of the Great Storm. As lightning struck the ancient oak tree in the village square, {protagonist} felt a strange tingling sensation. Suddenly, {protagonist} could control {magic} - a power not seen in the realm for centuries. But as {protagonist} was discovering these new abilities, darkness was spreading across {place}. The evil sorcerer Malakar had returned and was draining the life force from the land. With the guidance of a wise old wizard, {protagonist} learned to harness {magic}, gathering allies and growing stronger. In a final confrontation that lit up the sky, {protagonist} faced Malakar, using {magic} in ways never imagined to save {place} and restore balance to the world.",
                    
                    "{protagonist} had always felt different from the other inhabitants of {place}. Strange dreams haunted the nights, and unexplainable events occurred when emotions ran high. On the day of the centennial festival, as {protagonist} touched the sacred crystal in the town square, a surge of {magic} awakened within. The elders revealed a prophecy: {protagonist} was the chosen one who must save {place} from the shadow that emerges once every hundred years. With newfound abilities of {magic}, {protagonist} journeyed to the forbidden mountains, discovering ancient truths about {place}'s history and the true nature of the shadow. The final battle tested not just {protagonist}'s command of {magic}, but also wisdom, compassion, and the courage to make the ultimate sacrifice for {place}.",
                    
                    "The floating islands of {place} were held together by ancient magic, but the bridges between them had begun to crumble. {protagonist}, a student at the Academy of Mystic Arts, accidentally discovered an ability to control {magic} during a failed classroom experiment. This power, thought to be extinct, marked {protagonist} as the heir to the Founder's legacy. As {place} began to literally fall apart, {protagonist} had to master {magic} while evading the Void Seekers, who believed destroying {place} would usher in a new world order. Through forgotten ruins and secret libraries, {protagonist} pieced together the complex spell that could restore {place}, discovering that {magic} was actually the lifeforce that connected all the inhabitants of the floating realm."
                ]
            },
            "arabic": {
                "adventure": [
                    "في آثار {location} القديمة، اكتشف {hero} الشجاع خريطة تقود إلى {object}. لم تكن الرحلة سهلة، حيث كانت الفخاخ الخطيرة والمخلوقات الأسطورية تحرس الطريق. مسلحًا بالشجاعة والتصميم، غامر {hero} في المجهول. بعد التغلب على تحديات عديدة، بما في ذلك عبور نهر خطير وأبو الهول الذي يتحدث بالألغاز، وصل {hero} أخيرًا إلى الغرفة التي يُحفظ فيها {object}. لم يكن هذا الكنز قيمًا فحسب - بل كان يملك القدرة على تغيير العالم.",
                    
                    "كانت للقرية الصغيرة {location} أسطورة عن {object}، وهي قطعة أثرية غامضة يُقال إنها تمنح من يجدها قوى استثنائية. كان {hero}، وهو شخص فضولي ومغامر، مفتونًا دائمًا بهذه القصة. عندما بدأت أحداث غريبة تحدث في القرية، أدرك {hero} أن الوقت قد حان للبحث عن {object}. اختبرت الرحلة عبر الغابات المظلمة والجبال الضبابية كل مهارة يمتلكها {hero}. لكن مع كل عقبة تم التغلب عليها، أصبح {hero} أقوى وأكثر حكمة، ليكتشف في النهاية أن القوة الحقيقية لـ {object} لم تكن كما توقع أي شخص.",
                    
                    "قضى {hero} سنوات في دراسة النصوص القديمة حول {object}، المخفي في مكان ما في {location}. رفض العلماء هذه القصص باعتبارها أساطير، لكن {hero} كان يعرف الحقيقة. عندما ظهرت علامات مقلقة في السماء، تعرف عليها {hero} من النصوص - حان الوقت للعثور على {object}. كشفت البعثة إلى {location} عن مخاطر وعجائب تفوق الخيال. واجه {hero} ليس فقط التحديات الجسدية ولكن أيضًا المعضلات الأخلاقية التي شككت في كل ما اعتقده {hero}. في الغرفة الأخيرة، اكتشف {hero} أن {object} كان في الواقع بوابة إلى المعرفة التي من شأنها أن تغير إلى الأبد فهم البشرية للكون."
                ],
                "friendship": [
                    "في حي هادئ، عاشت قطة صغيرة تدعى {character1} وحدها، تبحث عن الطعام والمأوى. في يوم ممطر، اكتشفت {character1} روبوتًا مهجورًا يدعى {character2} في زقاق، متضررًا ومتوقفًا عن العمل. بدافع الفضول، بقيت {character1} في الجوار حتى قام مهندس لطيف بإعادة تنشيط {character2}. عندما عادت {character2} إلى الحياة، كانت مرتبكة وغير مستقرة. بدأت {character1}، رغم حذرها المبدئي، في زيارة {character2} يوميًا. ببطء، تشكلت صداقة غير عادية. كانت {character2} تحمي {character1} من كلاب الحي، بينما كانت {character1} ترشد {character2} عبر شوارع المدينة. علّم كل منهما الآخر عن وجهات نظرهما المختلفة جدًا في الحياة، مما يثبت أن الصداقة يمكن أن تتجاوز حتى أكبر الاختلافات.",
                    
                    "كانت {character1} أكثر الروبوتات تطوراً التي أنشأتها شركة تكنولوجيا المستقبل، مصممة للتعلم والتطور. ومع ذلك، شعرت {character1} بأن هناك شيئًا مفقودًا في وجودها المبرمج. في أحد الأيام، تسللت قطة ضالة تدعى {character2} إلى المختبر. بدلاً من طردها، قدمت لها {character1} الطعام. يومًا بعد يوم، عادت {character2}، وبدأت {character1} تشعر بعواطف لم تكن في برمجتها. عندما قررت الشركة إعادة ضبط خوارزميات التعلم الخاصة بـ {character1}، استشعرت {character2} الخطر بطريقة ما وخلقت تشتيتًا سمح لـ {character1} بالهروب. معًا، شرعا في مغامرة من شأنها إعادة تعريف معنى الحياة والاهتمام بكائن آخر.",
                    
                    "بعد عطل فني، هربت {character1}، وهي روبوت منزلي تجريبي، من منشأة الأبحاث واختبأت في مبنى مهجور. هناك، واجهت {character2}، وهي قطة ذكية في الشارع هسهست في البداية وخدشت المخلوق المعدني الغريب. ولكن عندما جاء الشتاء وانخفضت درجات الحرارة، استخدمت {character1} التدفئة الداخلية للحفاظ على دفء {character2}. تدريجيًا، بدأت {character2} تثق في الروبوت اللطيف. عندما عثر الباحثون أخيرًا على {character1}، ذهلوا عندما وجدوا أنها طورت استجابات عاطفية من خلال صداقتها مع {character2}. وبدلاً من استعادة الروبوت، قرروا دراسة هذه العلاقة الفريدة، مما أدى إلى تعلم دروس قيمة حول الارتباط والتعاطف من شأنها أن تغير علم الروبوتات إلى الأبد."
                ],
                "fantasy": [
                    "في مملكة {place} المسحورة، عاش {protagonist} الشاب حياة عادية حتى يوم العاصفة الكبرى. عندما ضرب البرق شجرة البلوط القديمة في ساحة القرية، شعر {protagonist} بإحساس غريب من الوخز. فجأة، أصبح {protagonist} قادرًا على التحكم في {magic} - وهي قوة لم تظهر في المملكة منذ قرون. ولكن بينما كان {protagonist} يكتشف هذه القدرات الجديدة، كان الظلام ينتشر في أنحاء {place}. لقد عاد الساحر الشرير مالاكار وكان يستنزف قوة الحياة من الأرض. بتوجيه من ساحر عجوز حكيم، تعلم {protagonist} ترويض {magic}، وجمع الحلفاء وأصبح أقوى. في مواجهة نهائية أضاءت السماء، واجه {protagonist} مالاكار، واستخدم {magic} بطرق لم تكن متخيلة من قبل لإنقاذ {place} واستعادة التوازن إلى العالم.",
                    
                    "شعر {protagonist} دائمًا بأنه مختلف عن سكان {place} الآخرين. طاردت الأحلام الغريبة الليالي، ووقعت أحداث غير قابلة للتفسير عندما كانت المشاعر مرتفعة. في يوم المهرجان المئوي، عندما لمس {protagonist} الكريستال المقدس في ساحة البلدة، استيقظت موجة من {magic} بداخله. كشف الشيوخ عن نبوءة: كان {protagonist} هو المختار الذي يجب أن ينقذ {place} من الظل الذي يظهر مرة واحدة كل مائة عام. مع القدرات الجديدة لـ {magic}، سافر {protagonist} إلى الجبال المحرمة، واكتشف حقائق قديمة حول تاريخ {place} والطبيعة الحقيقية للظل. اختبرت المعركة النهائية ليس فقط تحكم {protagonist} في {magic}، ولكن أيضًا الحكمة والتعاطف والشجاعة لتقديم التضحية النهائية من أجل {place}.",
                    
                    "كانت الجزر العائمة في {place} متماسكة بفضل السحر القديم، لكن الجسور بينها بدأت تتفكك. اكتشف {protagonist}، وهو طالب في أكاديمية الفنون الصوفية، عن طريق الصدفة قدرة على التحكم في {magic} خلال تجربة فاشلة في الفصل الدراسي. هذه القوة، التي كان يعتقد أنها انقرضت، وضعت علامة على {protagonist} كوريث لإرث المؤسس. مع بدء {place} في التفكك حرفيًا، كان على {protagonist} إتقان {magic} مع تجنب المتطلعين إلى الفراغ، الذين اعتقدوا أن تدمير {place} سيؤدي إلى نظام عالمي جديد. من خلال الأنقاض المنسية والمكتبات السرية، جمع {protagonist} أجزاء التعويذة المعقدة التي يمكن أن تعيد {place}، ليكتشف أن {magic} كان في الواقع قوة الحياة التي تربط جميع سكان المملكة العائمة."
                ]
            }
        }
    
    def is_safe_content(self, text: str) -> bool:
        """Check if content is safe to display (basic pattern filtering)"""
        # Check against banned patterns
        for pattern in self.banned_patterns:
            if re.search(pattern, text.lower()):
                logger.warning(f"Content filtered - banned pattern detected")
                return False
                
        # Check for repetitive patterns (sign of model failure)
        words = text.split()
        if len(words) > 20:
            unique_ratio = len(set(words)) / len(words)
            if unique_ratio < 0.4:  # Less than 40% unique words
                logger.warning(f"Content filtered - repetitive content (unique ratio: {unique_ratio})")
                return False
        
        return True
        
    def detect_language(self, text: str) -> str:
        """Detect if text is primarily Arabic or English"""
        arabic_char_count = sum(1 for char in text if '\u0600' <= char <= '\u06FF')
        return "arabic" if arabic_char_count > len(text) * 0.3 else "english"
        
    def analyze_sentiment(self, text: str, language: str) -> Dict[str, float]:
        """Analyze sentiment in the specified language"""
        try:
            if language == "arabic":
                if self.ar_sentiment_available:
                    result = self.ar_sentiment(text[:512])[0]
                    sentiment_label = result['label'] 
                    sentiment_score = float(result['score'])
                    
                    # Convert to standard format
                    if "positive" in sentiment_label.lower():
                        return {"positive": sentiment_score, "negative": 1-sentiment_score}
                    elif "negative" in sentiment_label.lower():
                        return {"negative": sentiment_score, "positive": 1-sentiment_score}
                    else:
                        return {"neutral": sentiment_score}
                else:
                    # Translate to English and use English sentiment
                    translated = self.translate(text, "arabic", "english")
                    return self.analyze_sentiment(translated, "english")
            else:
                # English sentiment analysis
                result = self.en_sentiment(text[:512])[0]
                sentiment_score = float(result['score'])
                
                if result['label'] == 'POSITIVE':
                    return {"positive": sentiment_score, "negative": 1-sentiment_score}
                else:
                    return {"negative": sentiment_score, "positive": 1-sentiment_score}
                
        except Exception as e:
            logger.error(f"Sentiment analysis error: {str(e)}")
            return {"neutral": 1.0}
            
    def translate(self, text: str, from_lang: str, to_lang: str) -> str:
        """Translate between languages"""
        try:
            if from_lang == to_lang:
                return text
                
            if from_lang == "arabic" and to_lang == "english":
                result = self.translator_ar_to_en(text[:512])
                return result[0]['translation_text']
            elif from_lang == "english" and to_lang == "arabic":
                if self.en_to_ar_available:
                    result = self.translator_en_to_ar(text[:512])
                    return result[0]['translation_text']
                else:
                    logger.warning("English to Arabic translation not available")
                    return text
            else:
                return text
        except Exception as e:
            logger.error(f"Translation error: {str(e)}")
            return text
            
    def _select_template(self, prompt: str, language: str) -> Tuple[str, Dict[str, str]]:
        """Select most appropriate template based on prompt keywords"""
        # Lower case for matching
        prompt_lower = prompt.lower()
        
        # Check for keywords to match
        template_type = "adventure"  # Default
        
        # Adventure keywords
        adventure_terms = ["adventure", "explorer", "ruins", "ancient", "discover", "treasure", "quest",
                          "مغامرة", "مستكشف", "آثار", "قديمة", "اكتشاف", "كنز"]
                          
        # Friendship keywords
        friendship_terms = ["friend", "friendship", "relationship", "together", "bond", "cat", "robot",
                           "صداقة", "صديق", "علاقة", "معًا", "رابطة", "قطة", "روبوت"]
                           
        # Fantasy keywords
        fantasy_terms = ["magic", "wizard", "spell", "kingdom", "power", "mystical", "enchanted",
                        "سحر", "ساحر", "تعويذة", "مملكة", "قوة", "غامض", "مسحور"]
        
        # Count matches
        adventure_matches = sum(1 for term in adventure_terms if term in prompt_lower)
        friendship_matches = sum(1 for term in friendship_terms if term in prompt_lower)
        fantasy_matches = sum(1 for term in fantasy_terms if term in prompt_lower)
        
        # Select template with most matches
        if friendship_matches > adventure_matches and friendship_matches > fantasy_matches:
            template_type = "friendship"
        elif fantasy_matches > adventure_matches and fantasy_matches > friendship_matches:
            template_type = "fantasy"
            
        # Extract potential parameters based on template type
        params = {}
        
        if template_type == "adventure":
            # Try to extract hero, object, location from prompt
            hero_match = re.search(r"(about|عن)\s+(?:a|an|the)?\s*(\w+)", prompt_lower)
            if hero_match:
                params["hero"] = hero_match.group(2).capitalize()
            else:
                params["hero"] = "the adventurer" if language == "english" else "المغامر"
                
            # Try to find object
            object_match = re.search(r"(find|discover|يجد|يكتشف)\s+(?:a|an|the)?\s*(\w+)", prompt_lower)
            if object_match:
                params["object"] = object_match.group(2)
            else:
                params["object"] = "lost treasure" if language == "english" else "الكنز المفقود"
                
            # Try to find location
            location_match = re.search(r"(in|at|في)\s+(?:a|an|the)?\s*(\w+)", prompt_lower)
            if location_match:
                params["location"] = location_match.group(2)
            else:
                params["location"] = "ancient temple" if language == "english" else "المعبد القديم"
                
        elif template_type == "friendship":
            # Handle friendship template parameters
            params["character1"] = "Whiskers" if language == "english" else "مشمش"
            params["character2"] = "Bolt" if language == "english" else "بولت"
            
            # Look for character mentions
            char_match = re.search(r"(between|بين)\s+(?:a|an|the)?\s*(\w+)", prompt_lower)
            if char_match:
                params["character1"] = char_match.group(2).capitalize()
                
        elif template_type == "fantasy":
            # Handle fantasy template parameters
            params["protagonist"] = "young wizard" if language == "english" else "الساحر الشاب"
            params["magic"] = "elemental magic" if language == "english" else "سحر العناصر"
            params["place"] = "mystical kingdom" if language == "english" else "المملكة السحرية"
            
            # Look for character mentions
            protag_match = re.search(r"(about|عن)\s+(?:a|an|the)?\s*(\w+)", prompt_lower)
            if protag_match:
                params["protagonist"] = protag_match.group(2).capitalize()
                
        return template_type, params
    
    def _customize_template(self, template: str, params: Dict[str, str]) -> str:
        """Fill template with parameters"""
        for key, value in params.items():
            placeholder = "{" + key + "}"
            template = template.replace(placeholder, value)
        return template
        
    def generate_story(
        self, 
        prompt: str, 
        language: str,
        template_name: Optional[str] = None,
        parameters: Optional[Dict[str, str]] = None
    ) -> Tuple[str, Dict[str, float]]:
        """Generate a story in the specified language with sentiment analysis"""
        if not self.initialized:
            error_msg = "Story generator was not properly initialized."
            return error_msg, {"neutral": 1.0}
            
        try:
            # Normalize language
            language = language.lower()
            if language not in ["english", "arabic"]:
                language = "english"
                
            # Detect input language and translate prompt if needed
            input_language = self.detect_language(prompt)
            if input_language != language:
                prompt = self.translate(prompt, input_language, language)
                
            # If template and parameters are provided, use them
            if template_name and parameters:
                if template_name in self.story_templates[language]:
                    # Get random template of the specified type
                    template_options = self.story_templates[language][template_name]
                    template = random.choice(template_options)
                    story = self._customize_template(template, parameters)
                else:
                    # Fallback to adventure if template not found
                    template = random.choice(self.story_templates[language]["adventure"])
                    story = self._customize_template(template, parameters)
            else:
                # Determine best template and parameters from prompt
                template_type, params = self._select_template(prompt, language)
                
                # Get random template of determined type
                template_options = self.story_templates[language][template_type]
                template = random.choice(template_options)
                
                # Fill template with extracted/default parameters
                story = self._customize_template(template, params)
                
            # Safety check
            if not self.is_safe_content(story):
                # Fallback to safe template
                template = random.choice(self.story_templates[language]["adventure"])
                default_params = {
                    "hero": "brave explorer" if language == "english" else "المستكشف الشجاع",
                    "object": "ancient artifact" if language == "english" else "القطعة الأثرية القديمة",
                    "location": "mysterious ruins" if language == "english" else "الآثار الغامضة"
                }
                story = self._customize_template(template, default_params)
                
            # Analyze sentiment
            emotions = self.analyze_sentiment(story, language)
            
            return story, emotions
            
        except Exception as e:
            logger.error(f"Story generation error: {str(e)}")
            # Return fallback story
            if language == "arabic":
                return "كان يا ما كان، في قديم الزمان، قصة لم تكتمل بعد...", {"neutral": 1.0}
            else:
                return "Once upon a time, there was a story waiting to be told...", {"neutral": 1.0}

class StoryManager:
    """Manages story templates and generation"""
    
    def __init__(self):
        self.templates = {
            "adventure": StoryTemplate(
                "Adventure Quest",
                "Tell a story about {hero} who must find {object} in {location}",
                "احكِ قصة عن {hero} الذي يجب أن يجد {object} في {location}",
                ["hero", "object", "location"]
            ),
            "friendship": StoryTemplate(
                "Friendship Tale",
                "Write about a friendship between {character1} and {character2}",
                "اكتب عن صداقة بين {character1} و {character2}",
                ["character1", "character2"]
            ),
            "fantasy": StoryTemplate(
                "Fantasy World",
                "Create a tale where {protagonist} discovers they have the power of {magic} and must save {place}",
                "اخلق قصة حيث يكتشف {protagonist} أن لديه قوة {magic} ويجب عليه إنقاذ {place}",
                ["protagonist", "magic", "place"]
            )
        }

class EmotionAnalyzer:
    """Legacy emotion analyzer class (kept for backward compatibility)"""
    
    def __init__(self):
        try:
            self.classifier = pipeline(
                "sentiment-analysis",
                model="distilbert-base-uncased-finetuned-sst-2-english",
                device="cpu"
            )
            logger.info("Initialized emotion analyzer successfully")
        except Exception as e:
            logger.error(f"Failed to initialize emotion analyzer: {str(e)}")
            self.classifier = None
        
    def analyze_emotions(self, text: str) -> Dict[str, float]:
        try:
            if not self.classifier:
                return {"neutral": 1.0}
                
            result = self.classifier(text[:512])[0]
            sentiment_score = float(result['score'])
            
            if result['label'] == 'POSITIVE':
                emotions = {
                    "positive": sentiment_score,
                    "negative": 1 - sentiment_score
                }
            else:
                emotions = {
                    "negative": sentiment_score,
                    "positive": 1 - sentiment_score
                }
            
            return emotions
            
        except Exception as e:
            logger.error(f"Emotion analysis error: {str(e)}")
            return {"neutral": 1.0}

class AIStoryteller:
    """Main class for story generation and management"""
    
    def __init__(self):
        self.story_manager = StoryManager()
        self.emotion_analyzer = EmotionAnalyzer()
        self.setup_models()

    def setup_models(self):
        """Initialize all required models and pipelines (on CPU by default)"""
        try:
            # Initialize the new multilingual story generator
            self.story_generator = MultilingualStoryGenerator()
            logger.info("Loaded multilingual story generator")
            
            # Translation Model (Arabic->English)
            translator_model_name = "Helsinki-NLP/opus-mt-ar-en"
            self.translator_tokenizer = MarianTokenizer.from_pretrained(translator_model_name)
            self.translator_model = MarianMTModel.from_pretrained(translator_model_name).to("cpu")
            logger.info("Loaded translation model")
    
            # Image Generation Model (Stable Diffusion), remains on CPU until needed
            self.pipe = StableDiffusionPipeline.from_pretrained(
                "runwayml/stable-diffusion-v1-5",
                torch_dtype=torch.float16
            )
            logger.info("Loaded Stable Diffusion pipeline (CPU -> GPU at runtime)")
    
        except Exception as e:
            logger.error(f"Error during model initialization: {str(e)}")
            raise

    def get_story_template(self, template_name: str, language: str) -> str:
        """Get template in specified language"""
        template = self.story_manager.templates.get(template_name)
        if not template:
            raise ValueError(f"Template {template_name} not found")
        return template.template_en if language.lower() == "english" else template.template_ar

    def fill_template(self, template_name: str, language: str, parameters: Dict[str, str]) -> str:
        """Fill template with provided parameters"""
        template = self.get_story_template(template_name, language)
        return template.format(**parameters)

    def translate_ar_to_en(self, text: str) -> str:
        """Translate Arabic text to English"""
        try:
            tokenized = self.translator_tokenizer([text], return_tensors="pt", truncation=True)
            translation = self.translator_model.generate(**tokenized)
            return self.translator_tokenizer.decode(translation[0], skip_special_tokens=True)
        except Exception as e:
            logger.error(f"Translation error: {str(e)}")
            raise ValueError(f"Failed to translate text: {str(e)}")

    def detect_language(self, text: str) -> str:
        """A simple heuristic to detect if text is mostly Arabic or English"""
        arabic_char_count = sum(1 for char in text if '\u0600' <= char <= '\u06FF')
        return "Arabic" if arabic_char_count > len(text) * 0.3 else "English"

    def translate_if_needed(self, text: str, from_lang: str, to_lang: str) -> str:
        """Legacy translate method (kept for backward compatibility)"""
        if from_lang == to_lang:
            return text
        
        if from_lang == "Arabic" and to_lang == "English":
            return self.translate_ar_to_en(text)
        elif from_lang == "English" and to_lang == "Arabic":
            logger.warning("English->Arabic translation not implemented.")
            return text
        return text

    def generate_text(
        self,
        prompt: str,
        language: str,
        template_name: Optional[str] = None,
        parameters: Optional[Dict[str, str]] = None
    ) -> str:
        """Generate text with optional template usage (CPU-based generation)"""
        try:
            # Convert language format to standard format for story generator
            lang = language.lower()
            
            # Generate story with sentiment analysis
            story, emotions = self.story_generator.generate_story(
                prompt, 
                lang,
                template_name, 
                parameters
            )
            
            # Format emotions for display
            if lang == "arabic":
                emotion_summary = "\n\nالنبرة العاطفية:\n"
                for emotion, score in emotions.items():
                    # Translate emotion names to Arabic
                    emotion_ar = {
                        "positive": "إيجابي",
                        "negative": "سلبي",
                        "neutral": "محايد"
                    }.get(emotion, emotion)
                    emotion_summary += f"- {emotion_ar}: {score:.1%}\n"
            else:
                emotion_summary = "\n\nEmotional tones:\n"
                for emotion, score in emotions.items():
                    emotion_summary += f"- {emotion}: {score:.1%}\n"
                    
            return story + emotion_summary
    
        except Exception as e:
            logger.error(f"Text generation error: {str(e)}")
            if language.lower() == "arabic":
                return f"عذراً، حدث خطأ أثناء إنشاء القصة: {str(e)}"
            else:
                return f"Error generating text: {str(e)}"

    @spaces.GPU
    def generate_story_and_scenes(
        self,
        prompt: str,
        language: str,
        template_name: Optional[str] = None,
        parameters: Optional[Dict[str, str]] = None,
        num_scenes: int = 3,
        style: str = "realistic"
    ) -> Tuple[str, List[Image.Image]]:
        """
        Single top-level function decorated with @spaces.GPU.
        1) Generate story on CPU.
        2) Move Stable Diffusion to GPU to generate multiple scenes.
        3) Move pipeline back to CPU at the end.
        """
        try:
            # Step 1: Generate story (CPU)
            story = self.generate_text(prompt, language, template_name, parameters)
            if "Error generating text" in story or "عذراً، حدث خطأ" in story:
                return story, []
            
            # Prepare story for image generation
            if language.lower() == "arabic":
                story_for_image = self.translate_ar_to_en(story)
            else:
                story_for_image = story

            # Step 2: Move pipe to GPU
            self.pipe = self.pipe.to("cuda")

            segments = [seg.strip() for seg in story_for_image.split('.') if seg.strip()]
            if len(segments) < num_scenes:
                selected_segments = segments
            else:
                step = max(1, len(segments) // num_scenes)
                selected_segments = segments[::step][:num_scenes]

            style_prompts = {
                "realistic": "photorealistic, highly detailed, 4k",
                "anime": "anime style, studio ghibli, detailed",
                "fantasy": "fantasy art, magical, detailed illustration"
            }
            
            scenes = []
            for segment in selected_segments:
                prompt_for_image = f"{segment}, {style_prompts.get(style, '')}"
                with torch.no_grad():
                    image = self.pipe(
                        prompt_for_image,
                        num_inference_steps=30,
                        guidance_scale=7.5
                    ).images[0]
                if image:
                    scenes.append(image)

            # Step 3: Move back to CPU
            self.pipe = self.pipe.to("cpu")
            torch.cuda.empty_cache()
            gc.collect()

            return story, scenes

        except Exception as e:
            logger.error(f"Story and scene generation error: {str(e)}")
            if hasattr(self, 'pipe'):
                self.pipe = self.pipe.to("cpu")
            torch.cuda.empty_cache()
            return f"Error generating story and scenes: {str(e)}", []

def build_gradio_interface() -> gr.Blocks:
    """Build and return the Gradio interface"""
    storyteller = AIStoryteller()

    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        # Header
        with gr.Row(variant="panel"):
            gr.Markdown(
                """
                # 📚 AI Bilingual Storyteller & Illustrator
                
                Create engaging stories in English or Arabic, with optional illustrations and emotional analysis.
                
                - 📝 **Basic Story**: Simple text generation
                - 🎯 **Template Story**: Guided story creation
                - 🎨 **Visual Story**: Story with scene illustrations
                """
            )

        # Tabs
        with gr.Tabs() as tabs:
            # Basic Story
            with gr.Tab("📝 Basic Story"):
                gr.Markdown("Enter a prompt, choose a language, and get a story with emotion analysis.")
                
                with gr.Row():
                    prompt_text = gr.Textbox(
                        label="Story Prompt",
                        placeholder="Once upon a time... / في يوم من الأيام...",
                        lines=3
                    )
                    language_choice = gr.Radio(
                        choices=["English", "Arabic"],
                        value="English",
                        label="Language"
                    )
                
                generate_button = gr.Button("✨ Generate Story")
                output_story = gr.Textbox(
                    label="Your Generated Story",
                    lines=8,
                    show_copy_button=True
                )
                
                generate_button.click(
                    fn=storyteller.generate_text,
                    inputs=[prompt_text, language_choice],
                    outputs=output_story
                )

            # Template Story
            with gr.Tab("🎯 Template Story"):
                gr.Markdown("Choose a template and fill parameters for a structured story.")
                
                with gr.Row():
                    template_choice = gr.Dropdown(
                        choices=list(storyteller.story_manager.templates.keys()),
                        label="Story Template",
                        value="adventure"
                    )
                    template_language = gr.Radio(
                        choices=["English", "Arabic"],
                        value="English",
                        label="Language"
                    )
                
                template_params = gr.JSON(
                    label="Template Parameters",
                    value={"hero": "brave knight", "object": "magical sword", "location": "enchanted forest"}
                )
                
                template_button = gr.Button("🎮 Generate from Template")
                template_output = gr.Textbox(
                    label="Your Generated Story",
                    lines=8,
                    show_copy_button=True
                )
                
                template_button.click(
                    fn=lambda t, l, p: storyteller.generate_text(
                        "", l, template_name=t, parameters=p
                    ),
                    inputs=[template_choice, template_language, template_params],
                    outputs=template_output
                )

            # Visual Story
            with gr.Tab("🎨 Visual Story"):
                gr.Markdown("Create a story with multiple illustrated scenes.")
                
                with gr.Row():
                    scene_prompt = gr.Textbox(
                        label="Story Prompt",
                        placeholder="Describe your story scene...",
                        lines=3
                    )
                    scene_language = gr.Radio(
                        choices=["English", "Arabic"],
                        value="English",
                        label="Language"
                    )
                
                with gr.Row():
                    num_scenes = gr.Slider(
                        minimum=1,
                        maximum=5,
                        value=3,
                        step=1,
                        label="Number of Scenes"
                    )
                    art_style = gr.Radio(
                        choices=["realistic", "anime", "fantasy"],
                        value="realistic",
                        label="Art Style",
                        interactive=True
                    )
                
                scene_button = gr.Button("🎭 Generate Story & Scenes")
                scene_story = gr.Textbox(
                    label="Your Generated Story",
                    lines=8,
                    show_copy_button=True
                )
                scene_gallery = gr.Gallery(
                    label="Story Scenes",
                    columns=3,
                    height="auto"
                )
                
                scene_button.click(
                    fn=storyteller.generate_story_and_scenes,
                    inputs=[
                        scene_prompt,
                        scene_language,
                        gr.Textbox(value=None, visible=False),
                        gr.Textbox(value=None, visible=False),
                        num_scenes,
                        art_style
                    ],
                    outputs=[scene_story, scene_gallery]
                )

        # Footer with Examples
        with gr.Row(variant="panel"):
            gr.Markdown("""
                ### Example Prompts
                
                **English:**
                - "Tell a story about a young wizard discovering a magical garden"
                - "Create an adventure about a brave explorer in ancient ruins"
                - "Write about a friendship between a cat and a robot"
                
                **Arabic:**
                - "احكِ قصة عن ساحر صغير يكتشف حديقة سحرية"
                - "اكتب مغامرة عن مستكشف شجاع في آثار قديمة"
                - "اكتب عن صداقة بين قطة وروبوت"
            """)

    return demo

if __name__ == "__main__":
    try:
        demo = build_gradio_interface()
        demo.launch()
    except Exception as e:
        logger.error(f"Application startup error: {str(e)}")
        raise