from __future__ import annotations

from typing import Any

from prompt_auto_lab.config.datasets import (
    DATASET_TO_EVAL_MAP,
    DEFAULT_EVAL_NAME,
    METRIC_OPT_DIRECTIONS,
)


def get_eval_key_metrics_from_dataset(dataset_name: str) -> list[str] | None:
    """Get key metrics name list by dataset name.

    Handles both legacy ``list[str]`` and new ``list[dict]`` formats.
    """
    mapping = DATASET_TO_EVAL_MAP.get(dataset_name)
    if not isinstance(mapping, dict):
        return None
    key_metrics = mapping.get("key_metrics")
    if not isinstance(key_metrics, list):
        return None
    names: list[str] = []
    for m in key_metrics:
        if isinstance(m, str) and m.strip():
            names.append(m.strip())
        elif isinstance(m, dict):
            name = m.get("name")
            if isinstance(name, str) and name.strip():
                names.append(name.strip())
    return names or None


def get_eval_name_from_dataset(dataset_name: str, default: str = DEFAULT_EVAL_NAME) -> str:
    """Get eval name by dataset name, with default fallback."""
    mapping = DATASET_TO_EVAL_MAP.get(dataset_name)
    if not isinstance(mapping, dict):
        return default
    eval_name = mapping.get("eval_name")
    if isinstance(eval_name, str) and eval_name.strip():
        return eval_name.strip()
    return default


def normalize_metrics_config(metrics_config: list[dict[str, Any]] | None) -> list[dict[str, Any]] | None:
    """Normalize metrics_config to UI-compatible schema."""
    if metrics_config is None or not isinstance(metrics_config, list):
        return None

    normalized: list[dict[str, Any]] = []
    for item in metrics_config:
        if not isinstance(item, dict):
            continue

        name_raw = item.get("name")
        name = str(name_raw).strip() if name_raw is not None else ""
        if not name:
            continue

        direction = str(item.get("optimization_direction", "higher_better")).strip()
        if direction not in METRIC_OPT_DIRECTIONS:
            direction = "higher_better"

        entry: dict[str, Any] = {
            "name": name,
            "is_llm_judged": bool(item.get("is_llm_judged", True)),
            "optimization_direction": direction,
        }

        desc = item.get("description")
        if isinstance(desc, str) and desc.strip():
            entry["description"] = desc.strip()

        purpose = item.get("measurement_purpose")
        if isinstance(purpose, str) and purpose.strip():
            entry["measurement_purpose"] = purpose.strip()

        normalized.append(entry)
    return normalized or None


def derive_key_metrics_from_metrics_config(metrics_config: list[dict[str, Any]] | None) -> list[str] | None:
    """Extract metric names from metrics_config."""
    if not metrics_config:
        return None
    key_metrics = [
        item.get("name")
        for item in metrics_config
        if isinstance(item, dict) and isinstance(item.get("name"), str) and item.get("name", "").strip()
    ]
    cleaned = [m.strip() for m in key_metrics if isinstance(m, str) and m.strip()]
    return cleaned or None


def build_default_metrics_config(key_metrics: list[str] | list[dict[str, Any]] | None) -> list[dict[str, Any]] | None:
    """Build default metrics_config using key_metrics.

    Accepts both legacy ``list[str]`` and new ``list[dict]`` formats.
    When dicts are provided, ``is_llm_judged``, ``optimization_direction``,
    ``description`` and ``measurement_purpose`` are preserved.
    """
    if not key_metrics:
        return None
    metrics_config: list[dict[str, Any]] = []
    for item in key_metrics:
        if isinstance(item, dict):
            name = item.get("name")
            if not isinstance(name, str) or not name.strip():
                continue
            entry: dict[str, Any] = {
                "name": name.strip(),
                "is_llm_judged": bool(item.get("is_llm_judged", True)),
                "optimization_direction": item.get("optimization_direction", "higher_better"),
            }
            desc = item.get("description")
            if isinstance(desc, str) and desc.strip():
                entry["description"] = desc.strip()
            purpose = item.get("measurement_purpose")
            if isinstance(purpose, str) and purpose.strip():
                entry["measurement_purpose"] = purpose.strip()
            metrics_config.append(entry)
        elif isinstance(item, str):
            name = item.strip()
            if not name:
                continue
            metrics_config.append(
                {
                    "name": name,
                    "is_llm_judged": True,
                    "optimization_direction": "higher_better",
                }
            )
    return metrics_config or None


def get_default_metrics_config_from_dataset(dataset_name: str) -> list[dict[str, Any]] | None:
    """Build default metrics_config from dataset mapping.

    Reads the raw ``key_metrics`` (list of dicts) directly from the MAP so
    that ``is_llm_judged``, ``optimization_direction``, ``description`` and
    ``measurement_purpose`` are preserved rather than falling back to defaults.
    """
    mapping = DATASET_TO_EVAL_MAP.get(dataset_name)
    if not isinstance(mapping, dict):
        return None
    raw_key_metrics = mapping.get("key_metrics")
    if not isinstance(raw_key_metrics, list):
        return None
    return build_default_metrics_config(raw_key_metrics)


def resolve_metrics_config(
    *,
    dataset_name: str,
    metrics_config: list[dict[str, Any]] | None,
) -> tuple[list[dict[str, Any]] | None, list[str] | None]:
    """
    Resolve normalized metrics_config and derived key_metrics.

    Priority:
    1) Provided metrics_config (normalized)
    2) Dataset default mapping
    """
    normalized_metrics_config = normalize_metrics_config(metrics_config)
    key_metrics = derive_key_metrics_from_metrics_config(normalized_metrics_config)
    if key_metrics is None:
        key_metrics = get_eval_key_metrics_from_dataset(dataset_name)
        if normalized_metrics_config is None:
            normalized_metrics_config = get_default_metrics_config_from_dataset(dataset_name)
    return normalized_metrics_config, key_metrics

