"""
PlumLoom Engine — LoRA Comparison Test
Genera imágenes de prueba para personajes con LoRAs duplicados.
Cada variante de LoRA se genera con la misma pose para comparación justa.
"""

import os
import json
import time
import requests
import sys
from pathlib import Path

# =============================================================================
# CONFIGURACIÓN
# =============================================================================

FORGE_URL = "http://127.0.0.1:7860"
OUTPUT_BASE = Path(r"C:\Users\Usuario\Downloads\NEW\Output\TEST")

# Parámetros de generación (consistentes para todos los tests)
GENERATION_PARAMS = {
    "negative_prompt": "lowres, bad anatomy, bad hands, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name, bad feet, multiple views, splits, 2girls, extra people",
    "width": 832,
    "height": 1216,
    "steps": 28,
    "cfg_scale": 6,
    "sampler_name": "DPM++ 2M Karras",
    "batch_size": 1,
    "n_iter": 1,
    "seed": -1,
    "clip_skip": 2,
}

# Pose estándar para todos los tests (Danbooru format)
TEST_POSE = "1girl, solo, front view, standing, smile, dynamic pose, dynamic shoot, simple background, white background"

# Master prompt base (sin personaje específico)
MASTER_PROMPT = "masterpiece, best quality, highly detailed, sharp focus, very aesthetic, very awa, detailed background, score_9, score_8_up, <lora:IFL_v1.0_IL:0.7> <lora:add-detail-xl:0.5> <lora:blackwhiplash(bwl)-guy90-Illust-Lorav3:0.6> blaash, <lora:scara2b_ill_v1.00:0.7> scara2b-style, <lora:Add_color_ILXLv1:3>, age up, mature female,"

# =============================================================================
# PERSONAJES CON LoRAs DUPLICADOS (para comparación)
# =============================================================================

LORA_TEST_CHARACTERS = {
    "Odette": {
        "variants": [
            {
                "name": "V1_Entrenamiento",
                "lora": "<lora:training_565609-20260702044914529:1>",
                "trigger": "odette",
                "appearance": "purple eyes, blue hair, swept bangs, sidelocks, hair ornament, braid, hair bun",
            },
            {
                "name": "V2_Genshin_Illustrious",
                "lora": "<lora:odette v0.91 illustrious:1>",
                "trigger": "odette (genshin impact)",
                "appearance": "blue hair, ahoge, purple eyes",
            },
            {
                "name": "V4_Genshin_Entrenamiento",
                "lora": "<lora:training_4293681-20260708121102976:1>",
                "trigger": "odette (genshin impact)",
                "appearance": "blue hair, purple eyes",
            },
        ],
    },
    "Yelan": {
        "variants": [
            {
                "name": "V1_Yelan",
                "lora": "<lora:Yelan:1>",
                "trigger": "Yelan",
                "appearance": "multicolored hair, black hair, blue hair, very short hair, green eyes, huge breasts, earrings, braid",
            },
            {
                "name": "V2_YelanGenshin_IXL",
                "lora": "<lora:YelanGenshin_IXL_v2:1>",
                "trigger": "zzYelan",
                "appearance": "green eyes, black hair, blue hair, multicolored hair, short hair, two-tone hair, jewelry, large breasts, bare shoulders, diagonal bangs, bob cut, cleavage",
            },
        ],
    },
    "Ellen_Joe": {
        "variants": [
            {
                "name": "V1_Illustrious",
                "lora": "<lora:Ellen Joe (Zenless Zone Zero) Illustrious:0.8>",
                "trigger": "Ellen Joe",
                "appearance": "red eyes, black hair, short hair, large breasts",
            },
            {
                "name": "V2_JN_Illustrious",
                "lora": "<lora:JN_Ellen_Joe_Illus:0.8>",
                "trigger": "Ellen Joe",
                "appearance": "hair ornament, hairclip, red eyes, black hair, short hair, colored inner hair, large breasts, fish tail, spiked hairband",
            },
        ],
    },
    "Beatrice": {
        "variants": [
            {
                "name": "V1_Beatrice",
                "lora": "<lora:Beatrice:1>",
                "trigger": "Beatrice from Re:Zero",
                "appearance": "long hair, blue eyes, blonde hair, twintails, hair ribbon, sidelocks, symbol-shaped pupils, drill hair, twin drills, small breast, tiny girl, petite",
            },
            {
                "name": "V2_se_beatrice",
                "lora": "<lora:se_beatrice_re_zero:1>",
                "trigger": "beatrice_re_zero",
                "appearance": "long hair, blue eyes, blonde hair, bow, ribbon, twin tails, hair ribbon, side locks, tights",
            },
        ],
    },
    "Hotori": {
        "variants": [
            {
                "name": "V1_hotori_nte",
                "lora": "<lora:hotori_nte:1>",
                "trigger": "hotori_nte",
                "appearance": "mature female, long hair, black hair, parted bangs, red eyes, hair ornament, earrings, makeup, large breasts",
            },
            {
                "name": "V2_Hotori_NTE",
                "lora": "<lora:Hotori_Neverness_to_Everness:1>",
                "trigger": "hotarinte",
                "appearance": "1girl, solo, long hair, black hair, red eyes, makeup",
            },
        ],
    },
    "Rebecca": {
        "variants": [
            {
                "name": "V1_RebeccaC77",
                "lora": "<lora:RebeccaC77IL:1>",
                "trigger": "reb77f",
                "appearance": "1girl, red sclera, colored sclera, twintails, green hair, colored skin, white skin, artificial eye",
            },
            {
                "name": "V2_Rebecca",
                "lora": "<lora:Rebecca:1>",
                "trigger": "rebecca",
                "appearance": "yellow eye, mechanical eye, red sclera, twintails, leg tattoo, colored skin, tattoo, petite, black hairband, green hair, neck tattoo, small breasts, headgear, long hair, stomach tattoo",
            },
        ],
    },
    "Rem": {
        "variants": [
            {
                "name": "V1_rem_basico",
                "lora": "<lora:rem:1>",
                "trigger": "rem",
                "appearance": "blue eyes, blue hair, short hair, hair over one eye, small breasts",
            },
            {
                "name": "V2_Rem_IlluXL",
                "lora": "<lora:ReZero_Rem_IlluXL:1>",
                "trigger": "rem",
                "appearance": "hair over one eye, blunt bangs, short hair, blue hair, blue eyes, dot nose, hair ribbon, purple ribbon, medium breasts",
            },
        ],
    },
    "Yixuan": {
        "variants": [
            {
                "name": "V1_Yixuan_IllustriousXL",
                "lora": "<lora:Yixuan_ZenlessZoneZero_Yixuan_IllustriousXL_v2:1>",
                "trigger": "yxn",
                "appearance": "large breasts, yellow eyes, white hair, ahoge, swept bangs, single hair intake, sidelocks, very long hair, half updo, hairclip, hair ornament",
            },
            {
                "name": "V2_yi_xuan_zzz",
                "lora": "<lora:yi_xuan_zzz_ilxl_goofy:1>",
                "trigger": "yi xuan (zenless zone zero)",
                "appearance": "yellow eyes, white hair, hairpin, hair intake, long hair",
            },
        ],
    },
    "Xianyun": {
        "variants": [
            {
                "name": "V1_Xianyun",
                "lora": "<lora:Xianyun:1>",
                "trigger": "Xianyun",
                "appearance": "navy blue hair, light green inner hair, very long ponytail, light green eyes, glasses, huge breasts, hairpin, hair ornaments, forehead",
            },
            {
                "name": "V2_Xianyun_Illustrious",
                "lora": "<lora:Xianyun_-_Genshin_Impact_-_Illustrious:1>",
                "trigger": "x1anyun",
                "appearance": "aqua eyes, green hair, black hair, very long hair, multicolored hair, bangs, ponytail",
            },
        ],
    },
    "Sense": {
        "variants": [
            {
                "name": "V1_Anime_ver",
                "lora": "<lora:sense-s1-illustriousxl-lora-nochekaiser:1>",
                "trigger": "sense, sense (sousou no frieren)",
                "appearance": "long hair, brown hair, hair between eyes, brown eyes, very long hair, ahoge, mature female, small breasts",
            },
            {
                "name": "V2_General_ver",
                "lora": "<lora:Sense illu:1>",
                "trigger": "sense",
                "appearance": "brown hair, hair between eyes, brown eyes, eyes visible through hair expressionless",
            },
        ],
    },
    "Ada_Wong": {
        "variants": [
            {
                "name": "V1_RE4_Remake",
                "lora": "<lora:NM_AW_re4remake:1>",
                "trigger": "NM_AW_re4remake",
                "appearance": "ada wong, ada, black hair, short hair, black eyes, medium breasts",
            },
            {
                "name": "V2_RE_Richy",
                "lora": "<lora:adawong-re-richy-v1_ixl:1>",
                "trigger": "adasweater",
                "appearance": "brown eyes, black hair, short hair, parted bangs, swept bangs, breasts, medium breasts, lips",
            },
        ],
    },
}

# =============================================================================
# FUNCIONES
# =============================================================================

def check_forge_connection():
    """Verifica que Forge Neo esté activo."""
    try:
        resp = requests.get(f"{FORGE_URL}/sdapi/v1/sd-models", timeout=10)
        if resp.status_code == 200:
            models = resp.json()
            print(f"[OK] Forge Neo conectado — {len(models)} modelos disponibles")
            return True
    except Exception as e:
        print(f"[ERROR] No se pudo conectar a Forge Neo: {e}")
    return False


def generate_image(prompt: str, negative_prompt: str, output_path: Path, seed: int = -1):
    """Genera una imagen con Forge Neo y la guarda."""
    payload = {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "width": GENERATION_PARAMS["width"],
        "height": GENERATION_PARAMS["height"],
        "steps": GENERATION_PARAMS["steps"],
        "cfg_scale": GENERATION_PARAMS["cfg_scale"],
        "sampler_name": GENERATION_PARAMS["sampler_name"],
        "batch_size": 1,
        "n_iter": 1,
        "seed": seed,
        "clip_skip": GENERATION_PARAMS["clip_skip"],
    }

    try:
        resp = requests.post(
            f"{FORGE_URL}/sdapi/v1/txt2img",
            json=payload,
            timeout=300,
        )
        if resp.status_code == 200:
            result = resp.json()
            import base64
            img_data = base64.b64decode(result["images"][0])
            output_path.parent.mkdir(parents=True, exist_ok=True)
            with open(output_path, "wb") as f:
                f.write(img_data)
            return True
        else:
            print(f"  [ERROR] Forge respondió {resp.status_code}: {resp.text[:200]}")
            return False
    except Exception as e:
        print(f"  [ERROR] Error generando: {e}")
        return False


def build_test_prompt(character_data: dict, variant: dict) -> str:
    """Construye el prompt de test para una variante."""
    parts = [
        MASTER_PROMPT,
        variant["lora"],
        variant["trigger"] + ",",
        variant.get("appearance", ""),
        "BREAK",
        TEST_POSE,
    ]
    return " ".join(p for p in parts if p)


def run_lora_tests():
    """Ejecuta todos los tests de comparación de LoRAs."""
    print("=" * 70)
    print("  PLUMLOOM ENGINE — LoRA COMPARISON TEST")
    print("=" * 70)
    print()

    if not check_forge_connection():
        print("\n[ABORTADO] Forge Neo no está disponible.")
        print("Inicia Forge Neo con: webui-user.bat --api")
        return

    total_tests = sum(len(c["variants"]) for c in LORA_TEST_CHARACTERS.values())
    completed = 0
    failed = 0

    print(f"\nTotal de tests a generar: {total_tests}")
    print(f"Personajes: {len(LORA_TEST_CHARACTERS)}")
    print()

    for char_name, char_data in LORA_TEST_CHARACTERS.items():
        char_folder = OUTPUT_BASE / char_name.replace(" ", "_")
        char_folder.mkdir(parents=True, exist_ok=True)

        print(f"\n{'-' * 50}")
        print(f"  {char_name} — {len(char_data['variants'])} variantes")
        print(f"  Carpeta: {char_folder}")
        print(f"{'-' * 50}")

        for variant in char_data["variants"]:
            variant_name = variant["name"]
            filename = f"{variant_name}.png"
            output_path = char_folder / filename

            if output_path.exists():
                print(f"  [SKIP] {filename} ya existe")
                completed += 1
                continue

            prompt = build_test_prompt(char_data, variant)
            print(f"\n  [GEN] {variant_name}")
            print(f"    LoRA: {variant['lora']}")
            print(f"    Trigger: {variant['trigger']}")

            success = generate_image(
                prompt=prompt,
                negative_prompt=GENERATION_PARAMS["negative_prompt"],
                output_path=output_path,
            )

            if success:
                print(f"    [OK] Guardado: {output_path.name}")
                completed += 1
            else:
                print(f"    [FAIL] No se pudo generar")
                failed += 1

            # Rate limit entre generaciones
            time.sleep(2)

    print(f"\n{'=' * 70}")
    print(f"  RESULTADOS: {completed} completados, {failed} fallidos")
    print(f"  Carpeta: {OUTPUT_BASE}")
    print(f"{'=' * 70}")


# =============================================================================
# MAIN
# =============================================================================

if __name__ == "__main__":
    run_lora_tests()
