"""
NeonChat ComfyUI Integration — Profilbild-Generierung für Characters
"""

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")))
CHARACTERS_DIR = Path(__file__).parent / "characters"
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_portrait_workflow(prompt: str, negative: str = "", width: int = 512, height: int = 512) -> dict:
    return {
        "3": {
            "class_type": "KSampler",
            "inputs": {
                "seed": int(time.time()) % (2**32),
                "steps": 25,
                "cfg": 7.0,
                "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": f"portrait, {prompt}, detailed face, high quality, masterpiece, best quality, 1girl, solo",
                "clip": ["4", 1]
            }
        },
        "7": {
            "class_type": "CLIPTextEncode",
            "inputs": {
                "text": f"lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, {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 generate_portrait(character_name: str, description: str) -> str | None:
    if not is_comfyui_running():
        return None
    
    AVATARS_DIR.mkdir(parents=True, exist_ok=True)
    
    prompt = f"anime style, {description}, simple background, headshot"
    negative = "blurry, deformed, ugly, realistic"
    
    workflow = build_portrait_workflow(prompt, negative, width=512, height=512)
    
    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
        
        safe_name = character_name.lower().replace(" ", "_")
        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