"""
NeonChat ComfyUI Integration — Profilbild-Generierung für Characters
Basiert auf dem Character-Aussehen (appearance)
"""

import json, os, time, shutil
from pathlib import Path
import httpx

COMFYUI_URL = os.environ.get("COMFYUI_URL", "http://localhost:8188")
COMFYUI_OUTPUT_DIR = Path(os.environ.get("COMFYUI_OUTPUT", str(Path.home() / "comfy" / "ComfyUI" / "output")))
AVATARS_DIR = Path(__file__).parent / "static" / "avatars"

def is_comfyui_running() -> bool:
    try:
        with httpx.Client(timeout=3) as client:
            resp = client.get(f"{COMFYUI_URL}/system_stats")
            return resp.status_code == 200
    except:
        return False

def build_context_from_messages(messages: list, character_name: str) -> str:
    """Extrahiert visuelle Hinweise aus den letzten Chat-Nachrichten."""
    if not messages:
        return ""
    
    # Nehme die letzten 10 Nachrichten
    recent = messages[-10:]
    
    # Sammle nur Bot-Nachrichten — die beschreiben was sie trägt/macht
    bot_msgs = [m["content"] for m in recent if m["role"] == "assistant"]
    if not bot_msgs:
        return ""
    
    # Keywords die das Aussehen verändern könnten
    outfit_keywords = {
        "umgezogen": "outfit change",
        "Kleid": "dress", "Jacke": "jacket", "Mantel": "coat", "Hoodie": "hoodie",
        "Bademant": "bathrobe", "Nachthemd": "nightgown", "Pyjama": "pajamas",
        "Bikini": "bikini", "Unterwäsch": "lingerie", "nackt": "nude",
        "Schwimmanzug": "swimsuit", "Sport": "sportswear", "Leggings": "leggings",
        "Krawatte": "necktie", "Anzug": "suit", "Uniform": "uniform",
        "Schürze": "apron", "Tank Top" : "tank top", "Crop-Top": "crop top",
        "Pullover": "sweater", "Bluse": "blouse", "Hemd": "shirt",
    }
    
    pose_keywords = {
        "lacht": "smiling", "weint": "crying", "lächelt": "smiling",
        "schläft": "sleeping", "schläfrig": "sleepy",
        "sitzt": "sitting", "steht": "standing", "liegt": "lying down",
        "fließt": "leaning forward", "geht": "walking",
        "schaut": "looking at viewer", "dreht": "turning away",
        "streckt": "stretching", "gähnt": "yawning",
        "wütend": "angry", "traurig": "sad", "überrascht": "surprised",
        "verlegen": "embarrassed", "errötet": "blushing",
    }
    
    scene_keywords = {
        "Café": "cafe interior", "Bar": "bar", "Küche": "kitchen",
        "Schlafzimmer": "bedroom", "Badezimmer": "bathroom",
        "Bett": "on bed", "Sofa": "on sofa", "Badewanne": "in bathtub",
        "Draußen": "outdoors", "Park": "park", "Straße": "street",
        "Regen": "rainy", "Sonne": "sunny", "Nacht": "night",
    }
    
    found_outfits = set()
    found_poses = set()
    found_scenes = set()
    
    combined = " ".join(bot_msgs[-3:])  # Letzte 3 Bot-Nachrichten
    
    for de, en in outfit_keywords.items():
        if de.lower() in combined.lower():
            found_outfits.add(en)
    for de, en in pose_keywords.items():
        if de.lower() in combined.lower():
            found_poses.add(en)
    for de, en in scene_keywords.items():
        if de.lower() in combined.lower():
            found_scenes.add(en)
    
    parts = []
    if found_outfits:
        parts.append(f"wearing {', '.join(found_outfits)}")
    if found_poses:
        parts.append(", ".join(found_poses))
    if found_scenes:
        parts.append(", ".join(found_scenes))
    
    return ", ".join(parts) if parts else ""


def build_prompt_from_appearance(character: dict, context: str = "") -> str:
    """Baut einen ComfyUI-Prompt aus dem Character-Aussehen."""
    app = character.get("appearance", {})
    
    style = app.get("style", "anime")
    gender = app.get("gender", "female")
    age = app.get("age", "25")
    height = app.get("height", "")
    build = app.get("build", "")
    hair_color = app.get("hair_color", "")
    hair_style = app.get("hair_style", "")
    eye_color = app.get("eye_color", "")
    skin = app.get("skin", "")
    clothing = app.get("clothing", "")
    distinctive = app.get("distinctive", "")
    
    # Style prefix
    # Style prefix and gender tags — adapt to style
    style_map = {
        "anime": ("anime style, anime art, ", True),
        "realistic": ("photorealistic, realistic, detailed skin texture, ", False),
        "comic": ("comic book style, western comic, ", False),
        "cartoon": ("cartoon style, ", False),
        "pixel": ("pixel art, ", False),
    }
    style_prefix, use_danbooru = style_map.get(style, ("anime style, ", True))
    
    # Gender — Danbooru tags for anime, natural language for others
    if use_danbooru:
        if gender == "female":
            gender_tag = "1girl, solo"
        elif gender == "male":
            gender_tag = "1boy, solo"
        else:
            gender_tag = "solo"
    else:
        if gender == "female":
            gender_tag = "young woman"
        elif gender == "male":
            gender_tag = "young man"
        else:
            gender_tag = "person"
    
    # Translation map: German appearance values → English tags for the model
    translate_map = {
        # Hair colors
        "dunkelbraun": "dark brown", "hellbraun": "light brown", "braun": "brown",
        "schwarz": "black", "blond": "blonde", "blonde": "blonde",
        "rot": "red", "rotbraun": "auburn", "weiß": "white", "grau": "gray",
        "blau": "blue", "grün": "green", "pink": "pink", "lila": "purple",
        "silber": "silver", "feuerrot": "fiery red", "platinum": "platinum blonde",
        # Hair styles
        "kurz": "short", "kurz, ungeordnet": "short messy", "lang": "long",
        "ungeordnet": "messy", "pony": "ponytail", "zöpfe": "twin braids",
        "hochsteckfrisur": "updo", "gelockt": "curly", "glatt": "straight",
        "wellig": "wavy", "kurz, gepflegt": "short neat", "schulterlang": "shoulder length",
        # Eye colors
        "braun": "brown", "blau": "blue", "grün": "green", "grau": "gray",
        "bernstein": "amber", "haselnuss": "hazel", "violett": "violet",
        # Skin
        "hell": "fair", "dunkel": "dark", "mittel": "medium", "oliv": "olive",
        "blass": "pale", "gebräunt": "tanned",
        # Build
        "schlank": "slim", "zierlich": "petite", "athletisch": "athletic",
        "muskulös": "muscular", "kurvig": "curvy", "kräftig": "stocky",
        "dünn": "thin", "groß": "tall", "klein": "short",
        # Clothing keywords
        "Café-Schürze über schwarzem T-Shirt": "cafe apron over black t-shirt",
        "Café-Schürze über schwarzem T-Shirt, Jeans": "cafe apron over black t-shirt, jeans",
        "Café-Schürze": "cafe apron", "Schürze": "apron", "T-Shirt": "t-shirt",
        "Jeans": "jeans", "Kleid": "dress", "Anzug": "suit", "Hoodie": "hoodie",
        "Jacke": "jacket", "Mantel": "coat", "Uniform": "uniform",
        "über": "over", "schwarzem": "black", "weißem": "white",
        "rotem": "red", "blauem": "blue", "grünem": "green",
        # Distinctive features
        "müde Augen": "tired eyes", "Augenringe": "dark circles under eyes",
        "Kaugummi": "chewing gum", "Brille": "glasses", "Tattoo": "tattoo",
        "Narben": "scars", "Sommersprossen": "freckles",
        "leichte Augenringe": "slight dark circles under eyes",
        "immer ein Kaugummi im Mund": "chewing gum",
    }
    
    def tr(val: str) -> str:
        """Übersetzt deutsche Wörter zu englischen Model-Tags."""
        if not val:
            return val
        result = val
        for de, en in sorted(translate_map.items(), key=lambda x: -len(x[0])):
            result = result.replace(de, en)
        return result
    
    # Build prompt — use established Danbooru tags for NoobAI-XL
    parts = [
        f"{style_prefix}{gender_tag}",
        f"{age} years old",
    ]
    if build:
        parts.append(tr(build))
    if hair_color and hair_style:
        parts.append(f"{tr(hair_color)} hair, {tr(hair_style)}")
    elif hair_color:
        parts.append(f"{tr(hair_color)} hair")
    elif hair_style:
        parts.append(tr(hair_style))
    if eye_color:
        parts.append(f"{tr(eye_color)} eyes")
    if skin:
        parts.append(f"{tr(skin)} skin")
    if clothing:
        parts.append(f"wearing {tr(clothing)}")
    if distinctive:
        parts.append(tr(distinctive))
    if context:
        parts.append(context)
    
    prompt = ", ".join(parts)
    prompt += ", looking at viewer, detailed face, masterpiece, best quality, very aesthetic, absurdres"
    
    return prompt

def build_portrait_workflow(prompt: str, negative: str = "", width: int = 1024, height: int = 1024) -> dict:
    return {
        "3": {
            "class_type": "KSampler",
            "inputs": {
                "seed": int(time.time()) % (2**32),
                "steps": 35,
                "cfg": 6.5,
                "sampler_name": "dpmpp_2m",
                "scheduler": "karras",
                "denoise": 1.0,
                "model": ["4", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "latent_image": ["5", 0]
            }
        },
        "4": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": "NoobAI-XL-v1.1.safetensors"}
        },
        "5": {
            "class_type": "EmptyLatentImage",
            "inputs": {"width": width, "height": height, "batch_size": 1}
        },
        "6": {
            "class_type": "CLIPTextEncode",
            "inputs": {
                "text": prompt,
                "clip": ["4", 1]
            }
        },
        "7": {
            "class_type": "CLIPTextEncode",
            "inputs": {
                "text": f"realistic, 3d, render, photorealistic, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, blurry, deformed, ugly, multiple views, split screen, repetition, duplicate, {negative}",
                "clip": ["4", 1]
            }
        },
        "8": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["3", 0], "vae": ["4", 2]}
        },
        "9": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": "neonchat_portrait", "images": ["8", 0]}
        }
    }

def queue_prompt(workflow: dict) -> str:
    with httpx.Client(timeout=10) as client:
        resp = client.post(f"{COMFYUI_URL}/prompt", json={"prompt": workflow})
        result = resp.json()
        return result.get("prompt_id", "")

def get_history(prompt_id: str) -> dict:
    try:
        with httpx.Client(timeout=5) as client:
            resp = client.get(f"{COMFYUI_URL}/history/{prompt_id}")
            return resp.json()
    except:
        return {}

def wait_for_image(prompt_id: str, timeout: int = 180) -> str | None:
    start = time.time()
    while time.time() - start < timeout:
        history = get_history(prompt_id)
        if prompt_id in history:
            outputs = history[prompt_id].get("outputs", {})
            if "9" in outputs:
                images = outputs["9"].get("images", [])
                if images:
                    return images[0].get("filename", "")
        time.sleep(3)
    return None

def build_i2i_workflow(prompt: str, negative: str, reference_image_path: str, width: int = 1024, height: int = 1024, denoise: float = 0.45) -> dict:
    """Image-to-Image Workflow — nutzt ein Referenzbild für Konsistenz."""
    return {
        "3": {
            "class_type": "KSampler",
            "inputs": {
                "seed": int(time.time()) % (2**32),
                "steps": 25,
                "cfg": 7.0,
                "sampler_name": "dpmpp_2m",
                "scheduler": "karras",
                "denoise": denoise,
                "model": ["4", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "latent_image": ["10", 0]
            }
        },
        "4": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": "NoobAI-XL-v1.1.safetensors"}
        },
        "5": {
            "class_type": "LoadImage",
            "inputs": {"image": reference_image_path}
        },
        "6": {
            "class_type": "CLIPTextEncode",
            "inputs": {
                "text": prompt,
                "clip": ["4", 1]
            }
        },
        "7": {
            "class_type": "CLIPTextEncode",
            "inputs": {
                "text": f"realistic, 3d, render, photorealistic, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, blurry, deformed, ugly, multiple views, split screen, repetition, duplicate, {negative}",
                "clip": ["4", 1]
            }
        },
        "8": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["3", 0], "vae": ["4", 2]}
        },
        "9": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": "neonchat_portrait", "images": ["8", 0]}
        },
        "10": {
            "class_type": "VAEEncode",
            "inputs": {"pixels": ["11", 0], "vae": ["4", 2]}
        },
        "11": {
            "class_type": "ImageScale",
            "inputs": {
                "image": ["5", 0],
                "upscale_method": "bilinear",
                "width": width,
                "height": height,
                "crop": "center"
            }
        }
    }

def upload_image_to_comfy(image_path: str) -> str:
    """Lädt ein Bild zu ComfyUI hoch und gibt den Dateinamen zurück."""
    import os
    filename = os.path.basename(image_path)
    with open(image_path, "rb") as f:
        files = {"image": (filename, f, "image/png")}
        with httpx.Client(timeout=30) as client:
            resp = client.post(f"{COMFYUI_URL}/upload/image", files=files)
            data = resp.json()
            return data.get("name", filename)

def generate_portrait(character_name: str, character: dict, context: str = "") -> str | None:
    """Generiert ein Profilbild. Nutzt Image-to-Image wenn ein Referenzbild existiert."""
    if not is_comfyui_running():
        return None
    
    AVATARS_DIR.mkdir(parents=True, exist_ok=True)
    
    prompt = build_prompt_from_appearance(character, context)
    style = character.get("appearance", {}).get("style", "anime")
    # Negative prompt depends on style
    if style == "anime":
        negative = "realistic, 3d, render, photorealistic, "
    else:
        negative = "anime, 3d, render, "
    negative += "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, blurry, deformed, ugly"
    
    # Check if we have a reference image for consistency
    safe_name = character_name.lower().replace(" ", "_")
    ref_image = AVATARS_DIR / f"{safe_name}.png"
    
    if ref_image.exists():
        # Image-to-Image for consistency
        # Higher denoise if context mentions outfit change
        denoise = 0.55 if any(kw in (context or "").lower() for kw in ["outfit change", "wearing ", "nude", "lingerie", "bikini", "dress", "hoodie", "coat", "jacket"]) else 0.45
        try:
            uploaded_name = upload_image_to_comfy(str(ref_image))
            workflow = build_i2i_workflow(prompt, negative, uploaded_name, 1024, 1024, denoise=denoise)
        except Exception as e:
            print(f"I2I upload failed, falling back to T2I: {e}")
            workflow = build_portrait_workflow(prompt, negative, 1024, 1024)
    else:
        # First generation — Text-to-Image
        workflow = build_portrait_workflow(prompt, negative, 1024, 1024)
    
    try:
        prompt_id = queue_prompt(workflow)
        if not prompt_id:
            return None
        
        filename = wait_for_image(prompt_id, timeout=180)
        if not filename:
            return None
        
        comfy_output = COMFYUI_OUTPUT_DIR / filename
        if not comfy_output.exists():
            return None
        
        dest = AVATARS_DIR / f"{safe_name}.png"
        shutil.copy2(str(comfy_output), str(dest))
        
        return f"/static/avatars/{safe_name}.png"
    except Exception as e:
        print(f"ComfyUI error: {e}")
        return None