from __future__ import annotations

import copy
import json
from pathlib import Path
from typing import Any, Mapping


def load_config(path: str | Path | None = None, overrides: Mapping[str, Any] | None = None) -> dict:
    """Load a Chimera JSON config and apply shallow dotted-key overrides."""
    if path is None:
        path = Path(__file__).resolve().parents[1] / "config.json"
    with open(path, "r", encoding="utf-8") as fh:
        cfg = json.load(fh)
    if overrides:
        cfg = copy.deepcopy(cfg)
        for key, value in overrides.items():
            cur = cfg
            parts = str(key).split(".")
            for part in parts[:-1]:
                cur = cur.setdefault(part, {})
            cur[parts[-1]] = value
    return cfg


def scale_config(config: dict, scale: str = "base") -> dict:
    """Return a safe CPU-scaled copy while preserving feature flags.

    The uploaded Chimera config targets a large model.  These presets keep all
    modules wired but resize dimensions so tests/fine-tuning fit commodity CPU
    memory (including 16 GB DDR5 machines).
    """
    cfg = copy.deepcopy(config)
    presets = {
        "nano": dict(hidden_size=128, intermediate_size=344, num_hidden_layers=4, num_heads=4, head_dim=32, vocab_size=min(cfg.get("vocab_size", 32000), 8192)),
        "tiny": dict(hidden_size=256, intermediate_size=688, num_hidden_layers=6, num_heads=4, head_dim=64, vocab_size=min(cfg.get("vocab_size", 32000), 32768)),
        "small": dict(hidden_size=512, intermediate_size=1376, num_hidden_layers=8, num_heads=8, head_dim=64, vocab_size=min(cfg.get("vocab_size", 32000), 65536)),
        "base": {},
    }
    if scale not in presets:
        raise ValueError(f"unknown scale {scale!r}; choose {sorted(presets)}")
    cfg.update(presets[scale])
    h = cfg["hidden_size"]
    cfg["num_heads"] = max(1, min(cfg.get("num_heads", 4), h // max(1, cfg.get("head_dim", 64))))
    cfg["head_dim"] = h // cfg["num_heads"]
    cfg.setdefault("backbone", {}).setdefault("moe", {})
    moe = cfg["backbone"]["moe"]
    moe["layers"] = [i for i in moe.get("layers", []) if i < cfg["num_hidden_layers"]]
    moe["n_routed_experts"] = min(int(moe.get("n_routed_experts", 4)), 4 if scale in {"nano", "tiny"} else 8)
    moe["n_shared_experts"] = min(int(moe.get("n_shared_experts", 1)), 1)
    moe["num_experts_per_tok"] = min(int(moe.get("num_experts_per_tok", 2)), moe["n_routed_experts"])
    moe["moe_intermediate_size"] = min(int(moe.get("moe_intermediate_size", h * 2)), max(64, cfg["intermediate_size"] // 2))
    loop = cfg.setdefault("looping", {})
    if cfg["num_hidden_layers"] < 8:
        loop["enabled"] = False
    else:
        loop["prelude"] = [0, min(1, cfg["num_hidden_layers"] - 1)]
        loop["loop"] = [2, max(2, cfg["num_hidden_layers"] - 3)]
        loop["coda"] = [max(0, cfg["num_hidden_layers"] - 2), cfg["num_hidden_layers"] - 1]
    cfg.setdefault("span_inference", {})["enabled"] = bool(cfg.get("span_inference", {}).get("enabled", True))
    return cfg


def tiny_config() -> dict:
    return scale_config(load_config(), "nano")