from __future__ import annotations

import json
from dataclasses import replace
from typing import Any

from .spec import ConfigPatch, Proposal, StudySpec, StudyState


TOPOLOGY_FLAG_KEYS = {
    "tensor-parallel-size",
    "data-parallel-size",
    "expert-parallel-size",
    "enable-expert-parallel",
}


def normalized_config_patch(config_patch: Any) -> dict[str, dict[str, Any]]:
    if isinstance(config_patch, ConfigPatch):
        env_patch: Any = config_patch.env_patch
        flag_patch: Any = config_patch.flag_patch
    elif isinstance(config_patch, dict):
        env_patch = config_patch.get("env_patch")
        flag_patch = config_patch.get("flag_patch")
    else:
        env_patch = {}
        flag_patch = {}
    return {
        "env_patch": _canonical_env_map(env_patch if isinstance(env_patch, dict) else {}),
        "flag_patch": _canonical_flag_map(flag_patch if isinstance(flag_patch, dict) else {}),
    }


def effective_config_signature(study: StudySpec, config_patch: Any) -> str:
    patch = normalized_config_patch(config_patch)
    payload = {
        "env": _canonical_env_map({**study.engine.base_envs, **patch["env_patch"]}),
        "flags": _canonical_flag_map({**study.engine.base_flags, **patch["flag_patch"]}),
    }
    return json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":"))


def materialize_proposal_for_execution(
    *,
    study: StudySpec,
    state: StudyState,
    proposal: Proposal,
) -> Proposal:
    flag_patch = dict(proposal.config_patch.flag_patch)
    env_patch = dict(proposal.config_patch.env_patch)
    if not flag_patch and not env_patch:
        return proposal
    if TOPOLOGY_FLAG_KEYS.intersection(flag_patch):
        return proposal
    if not state.best_trial_id:
        return proposal
    incumbent = next(
        (trial for trial in state.trials if trial.trial_id == state.best_trial_id),
        None,
    )
    if incumbent is None or not isinstance(incumbent.config_patch, dict):
        return proposal
    incumbent_patch = incumbent.config_patch.get("flag_patch")
    if not isinstance(incumbent_patch, dict):
        return proposal
    inherited_topology = {
        key: value
        for key, value in incumbent_patch.items()
        if key in TOPOLOGY_FLAG_KEYS and study.engine.base_flags.get(key) != value
    }
    if not inherited_topology:
        return proposal
    merged_flag_patch = dict(inherited_topology)
    merged_flag_patch.update(flag_patch)
    return replace(
        proposal,
        config_patch=ConfigPatch(
            env_patch=env_patch,
            flag_patch=merged_flag_patch,
        ),
    )


def materialized_effective_config_signature(
    *,
    study: StudySpec,
    state: StudyState,
    proposal: Proposal,
) -> str:
    materialized = materialize_proposal_for_execution(
        study=study,
        state=state,
        proposal=proposal,
    )
    return effective_config_signature(study, materialized.config_patch)


def tested_config_signature_index(study: StudySpec, state: StudyState) -> dict[str, list[str]]:
    index: dict[str, list[str]] = {}
    for trial in state.trials:
        signature = effective_config_signature(study, trial.config_patch)
        index.setdefault(signature, []).append(trial.trial_id)
    return index


def _canonical_env_map(payload: dict[str, Any]) -> dict[str, str]:
    return {str(key): str(value) for key, value in payload.items()}


def _canonical_flag_map(payload: dict[str, Any]) -> dict[str, Any]:
    return {str(key): _canonical_flag_value(value) for key, value in payload.items()}


def _canonical_flag_value(value: Any) -> Any:
    if value is None or isinstance(value, bool):
        return value
    if isinstance(value, int):
        return value
    if isinstance(value, float):
        return int(value) if value.is_integer() else value
    if isinstance(value, str):
        return _canonical_string_flag_value(value)
    if isinstance(value, list):
        return [_canonical_flag_value(item) for item in value]
    if isinstance(value, tuple):
        return [_canonical_flag_value(item) for item in value]
    if isinstance(value, dict):
        return {str(key): _canonical_flag_value(item) for key, item in value.items()}
    return str(value)


def _canonical_string_flag_value(value: str) -> Any:
    stripped = value.strip()
    if not stripped:
        return stripped
    try:
        parsed_int = int(stripped, 10)
    except ValueError:
        pass
    else:
        return parsed_int
    try:
        parsed_float = float(stripped)
    except ValueError:
        return stripped
    if parsed_float.is_integer():
        return int(parsed_float)
    return parsed_float
