"""
MCP Tool: get_model_features
Query metadata for a specified model, including supported task types, extended features, and default inference parameters.
"""

from .common import (
    _load_yaml,
    _MODEL_LIST_PATH,
    _MODEL_DEFAULTS_PATH,
    _IMAGE_GEN_FEATURES_PATH,
    _MODEL_ARCHITECTURES_PATH,
    _CHAIN_FEATURES_PATH,
    _TASK_DEFINITIONS,
)
from .error_schema import make_validation_error, make_not_found_error


def handle_get_model_features(model: str) -> dict:
    """Query metadata for a specified model: supported task types, extended features, and default inference parameters."""
    if not model:
        return make_validation_error(
            "Parameter 'model' is required.",
            missing_fields=["model"],
        )

    model_list = _load_yaml(_MODEL_LIST_PATH)
    model_defaults = _load_yaml(_MODEL_DEFAULTS_PATH)
    features_config = _load_yaml(_IMAGE_GEN_FEATURES_PATH)
    arch_config = _load_yaml(_MODEL_ARCHITECTURES_PATH)
    chain_features = _load_yaml(_CHAIN_FEATURES_PATH)

    found_arch = None
    checkpoints = model_list.get("Checkpoint", {})
    for arch_name, arch_data in checkpoints.items():
        if not isinstance(arch_data, dict):
            continue
        for m in arch_data.get("models", []):
            if m.get("display_name") == model:
                found_arch = arch_name
                break
        if found_arch:
            break

    if not found_arch:
        return make_not_found_error("model", model)

    architectures = arch_config.get("architectures", {})
    arch_info = architectures.get(found_arch, {})
    model_type = arch_info.get("model_type", found_arch.lower())

    arch_features = features_config.get(model_type, features_config.get("default", {}))
    enabled_chains = arch_features.get("enabled_chains", [])

    supported_features = []
    for chain_name in enabled_chains:
        if chain_name in chain_features:
            chain_data = chain_features[chain_name]
            visibility = chain_data.get("visibility", "public")
            if visibility == "public":
                supported_features.append(chain_name)
            else:
                generic_mapping = {
                    "krea2_controlnet": "controlnet",
                    "anima_controlnet_lllite": "controlnet",
                    "controlnet_model_patch": "controlnet",
                    "flux1_ipadapter": "ipadapter",
                    "sd3_ipadapter": "ipadapter",
                    "hidream_o1_reference": "reference_latent",
                }
                generic_name = generic_mapping.get(chain_name)
                if generic_name and generic_name not in supported_features:
                    supported_features.append(generic_name)

    arch_defaults_section = model_defaults.get(found_arch, {})
    arch_level_defaults = arch_defaults_section.get("_defaults", {})
    model_specific_defaults = arch_defaults_section.get(model, {})
    global_defaults = model_defaults.get("Default", {})

    merged_defaults = {**global_defaults, **arch_level_defaults, **model_specific_defaults}

    default_parameter = {
        "sampler": merged_defaults.get("sampler_name", "euler"),
        "scheduler": merged_defaults.get("scheduler", "simple"),
        "steps": merged_defaults.get("steps", 20),
        "cfg": merged_defaults.get("cfg", 1.0),
    }

    supported_tasks = [t["task_type"] for t in _TASK_DEFINITIONS]

    result = {
        "name": model,
        "model_architecture": found_arch,
        "supported_tasks": supported_tasks,
        "supported_features": supported_features,
        "default_parameter": default_parameter,
    }

    default_pos = model_specific_defaults.get(
        "positive_prompt",
        arch_level_defaults.get("positive_prompt", ""),
    )
    default_neg = model_specific_defaults.get(
        "negative_prompt",
        arch_level_defaults.get("negative_prompt", ""),
    )
    if default_pos:
        result["default_positive_prompt"] = default_pos
    if default_neg:
        result["default_negative_prompt"] = default_neg

    return result