"""Köppen-Geiger raster codes and the Beck et al. legend loader."""

from __future__ import annotations

import re
from pathlib import Path
from typing import Dict

# Beck et al legend key is expected as text file, but this default handles common codes.
DEFAULT_KOPPEN_CODE_MAP = {
    1: "Af",
    2: "Am",
    3: "Aw",
    4: "BWh",
    5: "BWk",
    6: "BSh",
    7: "BSk",
    8: "Csa",
    9: "Csb",
    10: "Csc",
    11: "Cwa",
    12: "Cwb",
    13: "Cwc",
    14: "Cfa",
    15: "Cfb",
    16: "Cfc",
    17: "Dsa",
    18: "Dsb",
    19: "Dsc",
    20: "Dsd",
    21: "Dwa",
    22: "Dwb",
    23: "Dwc",
    24: "Dwd",
    25: "Dfa",
    26: "Dfb",
    27: "Dfc",
    28: "Dfd",
    29: "ET",
    30: "EF",
}


def load_koppen_legend(legend_path: Path | None) -> Dict[int, str]:
    """Load Köppen raster codes, using defaults when no legend exists."""
    if legend_path is None:
        return DEFAULT_KOPPEN_CODE_MAP

    mapping: Dict[int, str] = {}
    for line in legend_path.read_text(encoding="utf-8").splitlines():
        text = line.strip()
        if not text or text.startswith("#"):
            continue
        # Handles patterns like:
        # "1:  Af ..." or "1 = Af" or "1 Af"
        match = re.match(r"^(\d+)\s*[:=]?\s*([A-Za-z]{2,3})\b", text)
        if not match:
            continue

        key = int(match.group(1))
        value = match.group(2)
        mapping[key] = value

    return mapping if mapping else DEFAULT_KOPPEN_CODE_MAP