from __future__ import annotations

import os
from pathlib import Path
from typing import Any, Dict

import yaml


def read_config(path: str | os.PathLike) -> Dict[str, Any]:
    with open(path, "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f)
    if cfg is None:
        cfg = {}
    cfg.setdefault("output_dir", "outputs")
    cfg.setdefault("missing_value", -99.0)
    cfg.setdefault("random_seed", 42)
    cfg.setdefault("overwrite", False)
    cfg.setdefault("save_predictions", True)
    cfg.setdefault("rmae_percent", True)
    cfg.setdefault("clip_precipitation_to_zero", True)
    cfg.setdefault("valid_regions", [1, 2, 3, 4, 5, 6, 7])
    cfg.setdefault("file_pattern", "lead{lead}_{source_model}_01_BL_Test_VN.nc")
    return cfg


def project_root_from_config(config_path: str | os.PathLike) -> Path:
    return Path(config_path).resolve().parents[1]


def resolve_path(path_value: str | os.PathLike, base_dir: str | os.PathLike | None = None) -> Path:
    p = Path(path_value)
    if p.is_absolute():
        return p
    if base_dir is None:
        return p.resolve()
    return (Path(base_dir) / p).resolve()


def prepare_output_dirs(cfg: Dict[str, Any], base_dir: str | os.PathLike | None = None) -> Dict[str, Path]:
    out = resolve_path(cfg.get("output_dir", "outputs"), base_dir)
    dirs = {
        "output": out,
        "tasks": out,
        "logs": out / "logs",
        "metrics": out / "metrics",
        "figures": out / "figures",
        "maps": out / "figures" / "maps",
        "corrected": out / "corrected",
        "done": out / "done",
    }
    for p in dirs.values():
        p.mkdir(parents=True, exist_ok=True)
    return dirs


def set_thread_env(n_threads: int = 1) -> None:
    # Avoid CPU oversubscription when many tasks run in parallel.
    for key in [
        "OMP_NUM_THREADS",
        "MKL_NUM_THREADS",
        "OPENBLAS_NUM_THREADS",
        "NUMEXPR_NUM_THREADS",
        "VECLIB_MAXIMUM_THREADS",
    ]:
        os.environ.setdefault(key, str(n_threads))