"""Tests para genconfig_to_sdcpp_args."""

import sys
import os
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))

from ml.genconfig_to_sdcpp_args import genconfig_to_sdcpp_args, _SAMPLER_MAP
from ml.generation_config import GenerationConfig


def _make_cfg(**overrides):
    """Crea un GenerationConfig sintetico para tests via model_validate / constructor."""
    defaults = dict(
        prompt="a cat",
        seed=1,
        steps=20,
        cfg_scale=7.0,
        sampler="euler",
        width=512,
        height=512,
        model={"name": "v1-5-pruned.ckpt", "model_type": "sd15", "path": "/models/v1-5.ckpt"},
    )
    defaults.update(overrides)
    # Normalizar loras a dicts si fueron pasados como LoraRef
    if "loras" in defaults:
        normalized = []
        for lr in defaults["loras"]:
            if hasattr(lr, "__dict__") and not isinstance(lr, dict):
                normalized.append({"path": lr.path, "weight": lr.weight, "scale": lr.scale})
            else:
                normalized.append(lr)
        defaults["loras"] = normalized
    try:
        return GenerationConfig.model_validate(defaults)
    except AttributeError:
        from ml.model_ref import ModelRef
        from ml.lora_ref import LoraRef
        m = defaults.pop("model")
        if isinstance(m, dict):
            m = ModelRef(**m)
        loras = defaults.pop("loras", [])
        built = [LoraRef(**lr) if isinstance(lr, dict) else lr for lr in loras]
        return GenerationConfig(model=m, loras=tuple(built), **defaults)


def _get_flag_value(args: list[str], flag: str) -> str | None:
    """Extrae el valor de un flag en la lista de args."""
    try:
        idx = args.index(flag)
        return args[idx + 1]
    except (ValueError, IndexError):
        return None


def _get_all_flag_values(args: list[str], flag: str) -> list[str]:
    """Extrae todos los valores de un flag repetido (ej. --lora)."""
    values = []
    for i, arg in enumerate(args):
        if arg == flag and i + 1 < len(args):
            values.append(args[i + 1])
    return values


class TestGenconfigToSdcppArgs(unittest.TestCase):

    def test_sampler_euler_a_se_mapea_a_euler_a_en_el_flag_sampling_method(self):
        """sampler euler_a se mapea a euler_a en el flag --sampling-method"""
        cfg = _make_cfg(sampler="euler_a")
        args = genconfig_to_sdcpp_args(cfg)
        self.assertEqual(_get_flag_value(args, "--sampling-method"), "euler_a")

    def test_sampler_dpm_pp_2m_se_mapea_a_dpmpp2m(self):
        """sampler dpm++2m se mapea a dpmpp2m"""
        cfg = _make_cfg(sampler="dpm++2m")
        args = genconfig_to_sdcpp_args(cfg)
        self.assertEqual(_get_flag_value(args, "--sampling-method"), "dpmpp2m")

    def test_lora_con_path_y_weight_se_agrega_como_lora_path_weight(self):
        """lora con path y weight se agrega como --lora path:weight"""
        cfg = _make_cfg(
            loras=[{"path": "/loras/detail.safetensors", "weight": 0.8}]
        )
        args = genconfig_to_sdcpp_args(cfg)
        lora_values = _get_all_flag_values(args, "--lora")
        self.assertEqual(len(lora_values), 1)
        self.assertEqual(lora_values[0], "/loras/detail.safetensors:0.8")

    def test_multiples_loras_generan_multiples_pares_lora(self):
        """multiples loras generan multiples pares --lora"""
        cfg = _make_cfg(
            loras=[
                {"path": "/loras/a.safetensors", "weight": 0.5},
                {"path": "/loras/b.safetensors", "weight": 1.0},
            ]
        )
        args = genconfig_to_sdcpp_args(cfg)
        lora_values = _get_all_flag_values(args, "--lora")
        self.assertEqual(len(lora_values), 2)
        self.assertIn("/loras/a.safetensors:0.5", lora_values)
        self.assertIn("/loras/b.safetensors:1.0", lora_values)

    def test_negative_prompt_none_produce_string_vacio_en_negative_prompt(self):
        """negative_prompt None produce string vacio en --negative-prompt"""
        cfg = _make_cfg(negative_prompt=None)
        args = genconfig_to_sdcpp_args(cfg)
        self.assertEqual(_get_flag_value(args, "--negative-prompt"), "")

    def test_model_path_tiene_prioridad_sobre_model_name_en_m(self):
        """model.path tiene prioridad sobre model.name en -m"""
        cfg = _make_cfg(
            model={"name": "hub-name", "model_type": "sd15", "path": "/local/path/model.ckpt"}
        )
        args = genconfig_to_sdcpp_args(cfg)
        self.assertEqual(_get_flag_value(args, "-m"), "/local/path/model.ckpt")

    def test_sin_path_usa_model_name_en_m(self):
        """sin path usa model.name en -m"""
        cfg = _make_cfg(
            model={"name": "runwayml/sd-v1-5", "model_type": "sd15", "path": None}
        )
        args = genconfig_to_sdcpp_args(cfg)
        self.assertEqual(_get_flag_value(args, "-m"), "runwayml/sd-v1-5")

    def test_args_contiene_flags_obligatorios(self):
        """args contiene --prompt --seed --steps --cfg-scale --sampling-method -W -H -m"""
        cfg = _make_cfg()
        args = genconfig_to_sdcpp_args(cfg)
        for flag in ["--prompt", "--seed", "--steps", "--cfg-scale", "--sampling-method", "-W", "-H", "-m"]:
            self.assertIn(flag, args, f"Flag faltante: {flag}")

    def test_sampler_map_cubre_todos_los_samplers_canonicos(self):
        """_SAMPLER_MAP cubre todos los samplers canonicos del dominio ml"""
        canonical = {"euler", "euler_a", "dpm++2m", "dpm++2m_v2", "heun", "dpm2", "lcm"}
        self.assertEqual(set(_SAMPLER_MAP.keys()), canonical)

    def test_seed_steps_width_height_se_convierten_a_string(self):
        """seed steps width height se convierten a string en los args"""
        cfg = _make_cfg(seed=42, steps=25, width=768, height=512)
        args = genconfig_to_sdcpp_args(cfg)
        self.assertEqual(_get_flag_value(args, "--seed"), "42")
        self.assertEqual(_get_flag_value(args, "--steps"), "25")
        self.assertEqual(_get_flag_value(args, "-W"), "768")
        self.assertEqual(_get_flag_value(args, "-H"), "512")


if __name__ == "__main__":
    unittest.main()
