import os
import sys
import random
import logging
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from api_client import ForgeClient, VisionClient
from prompt_builder import PromptBuilder
from vision_critic import VisionCritic

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUTPUT_DIR = os.path.join(BASE_DIR, "Output", "patreon_NSFW")
LOGS_DIR = os.path.join(BASE_DIR, "Logs")

os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(LOGS_DIR, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler(os.path.join(LOGS_DIR, "patreon_test.log"), encoding="utf-8"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

def main():
    import yaml
    config_path = os.path.join(BASE_DIR, "Config", "plumloom_engine.yaml")
    with open(config_path, "r", encoding="utf-8") as f:
        config = yaml.safe_load(f)

    forge_cfg = config["endpoints"]["forge_neo"]
    forge = ForgeClient(url=forge_cfg["url"], timeout=forge_cfg["timeout"])

    vision_cfg = config["endpoints"]["vision"]
    vision = VisionClient(
        url=vision_cfg["url"],
        model=vision_cfg["model"],
        timeout=vision_cfg["timeout"]
    )

    builder = PromptBuilder(BASE_DIR)

    print("=" * 55)
    print("  PLUMLOOM ENGINE v2 - PRUEBA PATREON NSFW")
    print("=" * 55)

    if not forge.check_status():
        print("[X] Forge Neo NO responde")
        return
    print("[OK] Forge Neo conectado")

    names = builder.get_character_names()
    print("[OK] {} personajes cargados".format(len(names)))

    char1 = random.choice(names)
    char2 = random.choice([n for n in names if n != char1])

    print("\nPersonajes: {} / {}".format(char1, char2))

    gen_cfg = config["generation"]
    critic = VisionCritic(vision_client=vision, max_retries=2, delay_between_requests=31)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    for idx, char_name in enumerate([char1, char2], 1):
        print("\n" + "=" * 55)
        print("  GENERANDO {}/2: {}".format(idx, char_name))
        print("=" * 55)

        prompt_pos, prompt_neg = builder.build_random(
            character_name=char_name,
            spice_level="explicit"
        )

        print("\nPrompt positivo:")
        print("  {}".format(prompt_pos))
        print("\nPrompt negativo:")
        print("  {}".format(prompt_neg))

        print("\nGenerando imagen ({}x{}, {} steps)...".format(
            gen_cfg['width'], gen_cfg['height'], gen_cfg['steps']))

        images, info = forge.txt2img(
            prompt=prompt_pos,
            negative_prompt=prompt_neg,
            width=gen_cfg["width"],
            height=gen_cfg["height"],
            steps=gen_cfg["steps"],
            cfg_scale=gen_cfg["cfg_scale"],
            sampler_name=gen_cfg["sampler_name"],
            clip_skip=gen_cfg["clip_skip"]
        )

        if not images:
            print("[X] Forge no devolvió imágenes")
            continue

        image_bytes = images[0]
        seed = info.get("seed", 0)
        print("[OK] Imagen generada (seed: {})".format(seed))

        print("\nEnviando a critico visual (OVHcloud Qwen2.5-VL)...")
        approved, verdict, details = critic.full_audit(image_bytes)

        if approved:
            filename = "{}_{}_seed{}.png".format(timestamp, char_name, seed)
            filepath = os.path.join(OUTPUT_DIR, filename)
            forge.save_image(image_bytes, filepath)
            print("[OK] APROBADA -> {}".format(filename))
        else:
            rejected_dir = os.path.join(BASE_DIR, "Output", "rejected")
            os.makedirs(rejected_dir, exist_ok=True)
            filename = "{}_{}_rej.png".format(timestamp, char_name)
            filepath = os.path.join(rejected_dir, filename)
            forge.save_image(image_bytes, filepath)
            print("[X] RECHAZADA: {}".format(verdict))

    print("\n" + "=" * 55)
    print("  PRUEBA PATREON COMPLETADA")
    print("  Carpeta: {}".format(OUTPUT_DIR))
    print("=" * 55)

if __name__ == "__main__":
    main()