prompt_engine.py

"""
Motor de procesamiento y mejora de prompts
"""

import re
from typing import Dict, List, Tuple, Optional
from model_configs import MODEL_CONFIGS, ENHANCEMENT_RULES, ARTIST_STYLES


class PromptAnalyzer:
    """Analiza prompts para identificar problemas y oportunidades de mejora"""
    
    def __init__(self):
        self.common_issues = {
            "muy_corto": lambda p: len(p.split()) < 5,
            "sin_estilo": lambda p: not any(word in p.lower() for word in 
                ["style", "estilo", "art", "photo", "painting", "render", "illustration"]),
            "sin_calidad": lambda p: not any(word in p.lower() for word in 
                ["quality", "detailed", "hd", "4k", "8k", "masterpiece", "professional"]),
            "sin_iluminacion": lambda p: not any(word in p.lower() for word in 
                ["light", "lighting", "luz", "iluminación", "bright", "dark", "shadow"]),
            "sin_composicion": lambda p: not any(word in p.lower() for word in 
                ["shot", "view", "angle", "close", "wide", "portrait", "landscape"]),
            "palabras_vagas": lambda p: any(word in p.lower() for word in 
                ["nice", "good", "beautiful", "cool", "bonito", "lindo", "bueno"])
        }
    
    def analyze(self, prompt: str) -> Dict:
        """Analiza el prompt y retorna un reporte"""
        issues = []
        suggestions = []
        score = 100
        
        for issue_name, check_func in self.common_issues.items():
            if check_func(prompt):
                issues.append(issue_name)
                score -= 15
        
        if "muy_corto" in issues:
            suggestions.append("🔸 Añade más detalles descriptivos al prompt")
        if "sin_estilo" in issues:
            suggestions.append("🎨 Especifica un estilo artístico (ej: 'oil painting', 'digital art')")
        if "sin_calidad" in issues:
            suggestions.append("✨ Añade modificadores de calidad (ej: 'highly detailed', '8k')")
        if "sin_iluminacion" in issues:
            suggestions.append("💡 Describe la iluminación (ej: 'dramatic lighting', 'golden hour')")
        if "sin_composicion" in issues:
            suggestions.append("📷 Especifica el encuadre o ángulo de cámara")
        if "palabras_vagas" in issues:
            suggestions.append("📝 Reemplaza palabras vagas por descripciones específicas")
        
        return {
            "score": max(0, score),
            "issues": issues,
            "suggestions": suggestions,
            "word_count": len(prompt.split()),
            "char_count": len(prompt)
        }


class PromptEnhancer:
    """Mejora prompts para modelos de generación de imágenes"""
    
    def __init__(self):
        self.analyzer = PromptAnalyzer()
        
    def translate_basic_terms(self, prompt: str) -> str:
        """Traduce términos básicos del español al inglés"""
        translations = {
            "mujer": "woman", "hombre": "man", "niño": "child", "niña": "girl",
            "perro": "dog", "gato": "cat", "casa": "house", "árbol": "tree",
            "montaña": "mountain", "mar": "sea", "cielo": "sky", "bosque": "forest",
            "ciudad": "city", "calle": "street", "grande": "large", "pequeño": "small",
            "hermoso": "beautiful", "hermosa": "beautiful", "antiguo": "ancient",
            "moderno": "modern", "oscuro": "dark", "brillante": "bright",
            "colorido": "colorful", "misterioso": "mysterious", "rojo": "red",
            "azul": "blue", "verde": "green", "amarillo": "yellow", "negro": "black",
            "blanco": "white", "dorado": "golden", "plateado": "silver",
            "realista": "realistic", "fantástico": "fantasy", "futurista": "futuristic",
            "retro": "retro", "minimalista": "minimalist"
        }
        
        result = prompt
        for es, en in translations.items():
            result = re.sub(rf'\b{es}\b', en, result, flags=re.IGNORECASE)
        
        return result
    
    def enhance_for_model(self, prompt: str, model: str, style: str = None, 
                          add_quality: bool = True, add_negative: bool = True) -> Dict:
        """Mejora el prompt para un modelo específico"""
        
        config = MODEL_CONFIGS.get(model, MODEL_CONFIGS["stable-diffusion-1.5"])
        
        enhanced = self.translate_basic_terms(prompt)
        
        if style and style in config["style_modifiers"]:
            style_tokens = config["style_modifiers"][style]
            enhanced = f"{enhanced}, {', '.join(style_tokens)}"
        
        if add_quality and config["quality_tokens"]:
            quality_str = ", ".join(config["quality_tokens"][:4])
            enhanced = f"{quality_str}, {enhanced}"
        
        if model == "dall-e-3":
            enhanced = self._format_for_dalle(enhanced)
        elif model == "flux":
            enhanced = self._format_for_flux(enhanced)
        elif model == "midjourney":
            enhanced = self._format_for_midjourney(enhanced, config)
        
        negative = ""
        if add_negative and config["negative_default"]:
            negative = ", ".join(config["negative_default"])
        
        analysis = self.analyzer.analyze(enhanced)
        
        return {
            "original": prompt,
            "enhanced": enhanced,
            "negative_prompt": negative,
            "model": config["name"],
            "analysis": analysis,
            "tokens_estimate": len(enhanced.split()),
            "max_tokens": config["max_tokens"]
        }
    
    def _format_for_dalle(self, prompt: str) -> str:
        """Formatea el prompt para DALL-E 3"""
        prompt = prompt.replace(", ", " with ")
        prompt = prompt.replace("  ", " ")
        
        if not any(prompt.lower().startswith(x) for x in ["a ", "an ", "the "]):
            prompt = f"A detailed image of {prompt}"
        
        return prompt
    
    def _format_for_flux(self, prompt: str) -> str:
        """Formatea el prompt para FLUX"""
        prompt = re.sub(r'\([^)]+:\d+\.?\d*\)', '', prompt)
        prompt = re.sub(r'\s+', ' ', prompt).strip()
        return prompt
    
    def _format_for_midjourney(self, prompt: str, config: Dict) -> str:
        """Formatea el prompt para Midjourney"""
        params = config.get("parameters", {})
        param_str = f" {params.get('version', '--v 6')} {params.get('quality', '--q 2')}"
        return f"{prompt}{param_str}"
    
    def suggest_variations(self, prompt: str, num_variations: int = 3) -> List[str]:
        """Genera variaciones del prompt"""
        variations = []
        
        art_styles = ["oil painting style", "digital art", "watercolor painting"]
        lighting = ["dramatic lighting", "soft natural light", "neon lights"]
        mood = ["ethereal atmosphere", "moody and dark", "vibrant and colorful"]
        
        variations.append(f"{prompt}, {art_styles[0]}, masterpiece")
        variations.append(f"{prompt}, {lighting[0]}, detailed")
        variations.append(f"{prompt}, {mood[0]}, high quality")
        
        return variations[:num_variations]
    
    def fix_common_errors(self, prompt: str) -> Tuple[str, List[str]]:
        """Corrige errores comunes en prompts"""
        fixes = []
        result = prompt
        
        if "  " in result:
            result = re.sub(r'\s+', ' ', result)
            fixes.append("Espacios múltiples eliminados")
        
        if ",," in result or ", ," in result:
            result = re.sub(r',\s*,', ',', result)
            fixes.append("Comas duplicadas corregidas")
        
        if "()" in result or "( )" in result:
            result = re.sub(r'\(\s*\)', '', result)
            fixes.append("Paréntesis vacíos eliminados")
        
        open_count = result.count('(')
        close_count = result.count(')')
        if open_count != close_count:
            if open_count > close_count:
                result += ')' * (open_count - close_count)
            else:
                result = '(' * (close_count - open_count) + result
            fixes.append("Paréntesis balanceados")
        
        problematic = ['|', '\\', '`', '~']
        for char in problematic:
            if char in result:
                result = result.replace(char, '')
                fixes.append(f"Carácter '{char}' eliminado")
        
        def fix_weight(match):
            weight = float(match.group(2))
            if weight < 0.1:
                weight = 0.1
            elif weight > 2.0:
                weight = 2.0
            return f"({match.group(1)}:{weight})"
        
        result = re.sub(r'\(([^:]+):(\d+\.?\d*)\)', fix_weight, result)
        
        return result.strip(), fixes


class PromptBuilder:
    """Constructor interactivo de prompts"""
    
    def __init__(self):
        self.reset()
    
    def reset(self):
        self.components = {
            "subject": "",
            "style": "",
            "quality": [],
            "lighting": "",
            "mood": "",
            "camera": "",
            "colors": [],
            "details": []
        }
        return self
    
    def set_subject(self, subject: str):
        self.components["subject"] = subject
        return self
    
    def set_style(self, style: str):
        if style in ARTIST_STYLES:
            self.components["style"] = ", ".join(ARTIST_STYLES[style])
        else:
            self.components["style"] = style
        return self
    
    def add_quality(self, quality_terms: List[str]):
        self.components["quality"].extend(quality_terms)
        return self
    
    def set_lighting(self, lighting: str):
        for category, data in ENHANCEMENT_RULES.items():
            if category == "lighting" and lighting in data["enhancements"]:
                self.components["lighting"] = ", ".join(data["enhancements"][lighting])
                break
        else:
            self.components["lighting"] = lighting
        return self
    
    def set_mood(self, mood: str):
        if "mood" in ENHANCEMENT_RULES and mood in ENHANCEMENT_RULES["mood"]["enhancements"]:
            self.components["mood"] = ", ".join(ENHANCEMENT_RULES["mood"]["enhancements"][mood])
        else:
            self.components["mood"] = mood
        return self
    
    def set_camera(self, camera: str):
        if "camera" in ENHANCEMENT_RULES and camera in ENHANCEMENT_RULES["camera"]["enhancements"]:
            self.components["camera"] = ", ".join(ENHANCEMENT_RULES["camera"]["enhancements"][camera])
        else:
            self.components["camera"] = camera
        return self
    
    def add_colors(self, colors: List[str]):
        self.components["colors"].extend(colors)
        return self
    
    def add_details(self, details: List[str]):
        self.components["details"].extend(details)
        return self
    
    def build(self) -> str:
        """Construye el prompt final"""
        parts = []
        
        if self.components["quality"]:
            parts.append(", ".join(self.components["quality"]))
        
        if self.components["subject"]:
            parts.append(self.components["subject"])
        
        if self.components["style"]:
            parts.append(self.components["style"])
        
        if self.components["colors"]:
            parts.append(", ".join(self.components["colors"]))
        
        if self.components["lighting"]:
            parts.append(self.components["lighting"])
        
        if self.components["mood"]:
            parts.append(self.components["mood"])
        
        if self.components["camera"]:
            parts.append(self.components["camera"])
        
        if self.components["details"]:
            parts.append(", ".join(self.components["details"]))
        
        return ", ".join(filter(None, parts))