"""Shared configuration loading and resolution for the acdc CLI and services.

This is the single source of truth for reading ``config/deploy_config.yaml``,
resolving environment and study settings, and locating the repository root. It
replaces the ``load_config()`` / config-extraction logic that was previously
copy-pasted across the ``services/`` scripts.

Design notes
------------
* ``resolve_env`` / ``resolve_study`` take an already-loaded ``cfg`` dict so they
  are trivially unit-testable without touching the filesystem.
* Study keys in ``deploy_config.yaml`` embed the environment
  (e.g. ``ausdiab_staging``). ``resolve_study`` lets a caller pass the bare study
  (``ausdiab``) plus ``--env staging`` and derives ``ausdiab_staging``, while still
  accepting the literal key for back-compat.
* The repo root is derived from this file's location so nothing depends on a
  machine-specific absolute path (the old ``common.argo_script_dir`` hazard).
"""
from __future__ import annotations

import functools
import os
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Tuple

import yaml

DEFAULT_CONFIG_RELPATH = "config/deploy_config.yaml"
DEFAULT_REGION = "ap-southeast-2"

#: Where the data dictionary is fetched from. The download URL is composed as
#: ``{base_url}/{repo}/refs/tags/{version}/{schema_path}``. Set it globally under
#: ``common.dictionary`` (keys ``base_url``/``repo``/``schema_path``), or override
#: it for one environment with the flat keys ``dictionary_base_url`` /
#: ``dictionary_repo`` / ``dictionary_schema_path`` in that env's block.
#:
#: The nested-vs-flat asymmetry is deliberate: env blocks in deploy_config.yaml
#: are single-level, and ``set_env_value`` can only write single-level keys.
#:
#: These defaults exist so a config written before the dictionary source became
#: configurable keeps resolving to the same URL it always did.
DEFAULT_DICT_BASE_URL = "https://raw.githubusercontent.com"
DEFAULT_DICT_REPO = "AustralianBioCommons/acdc-schema-json"
DEFAULT_DICT_SCHEMA_PATH = "dictionary/prod_dict/acdc_schema.json"

#: Repo-relative directory that ``services/dictionary/pull_dict.sh`` downloads into.
DICT_SCHEMA_DIR_RELPATH = "services/dictionary/schemas"

#: Keys every environment profile block must define.
REQUIRED_ENV_KEYS = (
    "dictionary_version",
    "aws_secret_name",
    "schema_s3_uri",
    "domain",
    "app_name",
    "namespace",
    "cluster_name",
)

#: Env-profile keys that ``set_env_value`` may write. REQUIRED_ENV_KEYS plus the
#: optional block keys; excludes derived fields (``name``/``is_ec2``) and
#: ``region``, which lives under ``common.aws.region``, not per-env.
SETTABLE_ENV_KEYS = REQUIRED_ENV_KEYS + (
    "aws_profile",
    "eks_arn",
    "ec2_instance_id",
    "ssh_key",
    "ssh_user",
    # Per-env overrides of the dictionary source. The global values live under
    # `common.dictionary`, which set_env_value cannot write (it only touches env
    # blocks) -- change those by editing the YAML directly.
    "dictionary_base_url",
    "dictionary_repo",
    "dictionary_schema_path",
)

#: Top-level keys in the config that are NOT environment profiles.
_NON_ENV_KEYS = {"common", "studies"}

#: Suffixes used to strip the environment from a study key.
_STUDY_ENV_SUFFIXES = ("_staging", "_prod", "_test")


class ConfigError(KeyError):
    """Raised when an env/study cannot be resolved or required keys are missing.

    Subclasses ``KeyError`` so existing ``except KeyError`` handlers still catch
    it, but carries a human-readable message that doubles as CLI help.
    """

    def __str__(self) -> str:  # KeyError repr adds quotes; we want the raw text
        return self.args[0] if self.args else ""


def repo_root() -> Path:
    """Return the repository root directory.

    Resolved relative to this file (``src/acdc_aws_etl_pipeline/config.py`` →
    three parents up). Override with ``$ACDC_REPO_ROOT`` for unusual layouts.
    """
    override = os.getenv("ACDC_REPO_ROOT")
    if override:
        return Path(override)
    return Path(__file__).resolve().parents[2]


def find_config_path(override: Optional[str] = None) -> Path:
    """Locate ``deploy_config.yaml``.

    Precedence: explicit ``override`` > ``$ACDC_CONFIG`` > repo-root default.
    """
    if override:
        return Path(override)
    env = os.getenv("ACDC_CONFIG")
    if env:
        return Path(env)
    return repo_root() / DEFAULT_CONFIG_RELPATH


@functools.lru_cache(maxsize=None)
def _load_yaml_cached(path_str: str) -> dict:
    with open(path_str, "r", encoding="utf-8") as f:
        return yaml.safe_load(f) or {}


def load_config(path: Optional[str] = None) -> dict:
    """Load ``deploy_config.yaml`` (cached by resolved path)."""
    return _load_yaml_cached(str(find_config_path(path)))


def set_env_value(
    env: str, key: str, value: str, path: Optional[str] = None
) -> Tuple[Optional[str], str]:
    """Set one key in an environment profile and write the config back to disk.

    Uses ruamel.yaml in round-trip mode (reads stay on PyYAML) so comments, key
    order, and formatting in the hand-maintained YAML survive the rewrite.
    Returns ``(old_value, new_value)``; ``old_value`` is ``None`` if the key was
    not previously set.
    """
    # Imported here so the (read-only) rest of this module keeps zero
    # dependency on ruamel.
    from ruamel.yaml import YAML

    cfg_path = find_config_path(path)
    yaml_rt = YAML()
    yaml_rt.preserve_quotes = True
    # Keep explicit nulls explicit: ruamel dumps None as an empty scalar by
    # default, which would rewrite every untouched `key: null` line.
    yaml_rt.representer.add_representer(
        type(None),
        lambda rep, _: rep.represent_scalar("tag:yaml.org,2002:null", "null"),
    )
    with open(cfg_path, "r", encoding="utf-8") as f:
        doc = yaml_rt.load(f) or {}

    if env not in doc or env in _NON_ENV_KEYS:
        raise ConfigError(
            f"Environment '{env}' not found in config. "
            f"Valid environments: {', '.join(list_envs(doc))}"
        )
    if key not in SETTABLE_ENV_KEYS:
        raise ConfigError(
            f"Key '{key}' is not a settable environment key. "
            f"Settable keys: {', '.join(SETTABLE_ENV_KEYS)}"
        )

    old = doc[env].get(key)
    doc[env][key] = value
    with open(cfg_path, "w", encoding="utf-8") as f:
        yaml_rt.dump(doc, f)

    # The PyYAML reader caches by path for the process lifetime; without this a
    # subsequent load_config() in the same process would return the stale,
    # pre-write dict.
    _load_yaml_cached.cache_clear()
    return old, value


def env_base(env: str) -> str:
    """Strip a trailing ``_ec2`` suffix: ``staging_ec2`` -> ``staging``."""
    return env[:-4] if env.endswith("_ec2") else env


def aws_region(cfg: dict) -> str:
    return cfg.get("common", {}).get("aws", {}).get("region", DEFAULT_REGION)


#: Maps a ``common.dictionary`` key to its flat env-block equivalent and default.
_DICT_PARTS = {
    "base_url": ("dictionary_base_url", DEFAULT_DICT_BASE_URL),
    "repo": ("dictionary_repo", DEFAULT_DICT_REPO),
    "schema_path": ("dictionary_schema_path", DEFAULT_DICT_SCHEMA_PATH),
}


def dictionary_setting(cfg: dict, env: Optional[str], part: str) -> str:
    """Resolve one part of the dictionary source location for ``env``.

    Precedence: ``<env>.dictionary_<part>`` > ``common.dictionary.<part>`` > the
    module ``DEFAULT_DICT_*`` constant. Like :func:`aws_region` but with an extra
    per-env override layer.

    An explicit ``null`` or empty value is treated as unset, so a half-filled
    block falls back to the default instead of composing a broken URL. The
    ``or {}`` guards matter: a bare ``common:`` or ``dictionary:`` line in YAML
    parses to ``None``, which would make ``.get()`` chaining raise.
    """
    env_key, default = _DICT_PARTS[part]
    block = (cfg.get(env) or {}) if env else {}
    value = block.get(env_key)
    if not value:
        value = ((cfg.get("common") or {}).get("dictionary") or {}).get(part)
    return value or default


def list_envs(cfg: dict) -> List[str]:
    """Return the environment profile names (excludes ``common``/``studies``)."""
    return sorted(k for k in cfg if k not in _NON_ENV_KEYS)


def list_studies(cfg: dict) -> List[str]:
    """Return unique bare study names (env suffix stripped) for friendly help."""
    prefixes = set()
    for key in cfg.get("studies", {}):
        for suffix in _STUDY_ENV_SUFFIXES:
            if key.endswith(suffix):
                prefixes.add(key[: -len(suffix)])
                break
        else:
            prefixes.add(key)
    return sorted(prefixes)


def argo_script_dir(cfg: dict) -> Path:
    """Return the directory holding the ArgoCD/k8s shell scripts.

    Derived from the repo root so it is machine-independent. Falls back to a
    configured ``common.argo_script_dir`` only if the derived path is missing.
    """
    derived = repo_root() / "services" / "k8s_ops"
    if derived.is_dir():
        return derived
    configured = cfg.get("common", {}).get("argo_script_dir")
    return Path(configured) if configured else derived


@dataclass(frozen=True)
class EnvConfig:
    """Resolved settings for one environment profile."""

    name: str
    is_ec2: bool
    region: str
    dictionary_version: str
    aws_profile: Optional[str]
    aws_secret_name: str
    schema_s3_uri: str
    domain: str
    app_name: str
    namespace: str
    cluster_name: str
    eks_arn: Optional[str] = None
    ec2_instance_id: Optional[str] = None
    ssh_key: Optional[str] = None
    ssh_user: Optional[str] = None
    # Where this env's dictionary is fetched from; see dictionary_setting. Real
    # defaults (not None) so a hand-built EnvConfig still composes a valid URL.
    dictionary_base_url: str = DEFAULT_DICT_BASE_URL
    dictionary_repo: str = DEFAULT_DICT_REPO
    dictionary_schema_path: str = DEFAULT_DICT_SCHEMA_PATH


def resolve_env(cfg: dict, env: str) -> EnvConfig:
    """Resolve and validate an environment profile from the config dict."""
    block = cfg.get(env)
    if block is None:
        raise ConfigError(
            f"Environment '{env}' not found in config. "
            f"Valid environments: {', '.join(list_envs(cfg))}"
        )
    missing = [k for k in REQUIRED_ENV_KEYS if k not in block]
    if missing:
        raise ConfigError(
            f"Environment '{env}' is missing required key(s): "
            f"{', '.join(missing)}"
        )
    return EnvConfig(
        name=env,
        is_ec2=env.endswith("_ec2"),
        region=aws_region(cfg),
        dictionary_version=block["dictionary_version"],
        aws_profile=block.get("aws_profile"),
        aws_secret_name=block["aws_secret_name"],
        schema_s3_uri=block["schema_s3_uri"],
        domain=block["domain"],
        app_name=block["app_name"],
        namespace=block["namespace"],
        cluster_name=block["cluster_name"],
        eks_arn=block.get("eks_arn"),
        ec2_instance_id=block.get("ec2_instance_id"),
        ssh_key=block.get("ssh_key"),
        ssh_user=block.get("ssh_user"),
        dictionary_base_url=dictionary_setting(cfg, env, "base_url"),
        dictionary_repo=dictionary_setting(cfg, env, "repo"),
        dictionary_schema_path=dictionary_setting(cfg, env, "schema_path"),
    )


def _join_url(*parts: str) -> str:
    """Join URL segments with exactly one ``/`` between them.

    Each segment is stripped of leading/trailing slashes, so a trailing ``/`` on
    ``base_url`` or a leading ``/`` on ``schema_path`` in the YAML cannot produce
    a ``//`` mid-path (which raw.githubusercontent.com 404s on). Only the *ends*
    of each segment are touched, so the ``//`` in ``https://`` survives.
    """
    cleaned = [str(p).strip().strip("/") for p in parts]
    return "/".join(p for p in cleaned if p)


def dictionary_url(env_cfg: EnvConfig, version: Optional[str] = None) -> str:
    """Build the download URL for an environment's data dictionary.

    Composed as ``{base_url}/{repo}/refs/tags/{version}/{schema_path}`` from the
    env's resolved dictionary source. ``version`` overrides the env's
    ``dictionary_version`` (this backs ``--version`` on the CLI).
    """
    return _join_url(
        env_cfg.dictionary_base_url,
        env_cfg.dictionary_repo,
        "refs/tags",
        version or env_cfg.dictionary_version,
        env_cfg.dictionary_schema_path,
    )


def dictionary_filename(env_cfg: EnvConfig, version: Optional[str] = None) -> str:
    """Local basename for a pulled dictionary, e.g. ``acdc_schema_v1.1.6.json``.

    Derived from the configured ``schema_path``'s basename with the version
    inserted before the extension -- the same rule ``pull_dict.sh`` applies. The
    CLI passes this to that script as its optional second argument, so the
    downloaded name is decided in one place instead of being re-derived from the
    URL by a regex.
    """
    v = version or env_cfg.dictionary_version
    # `##*/` also strips any directory component, so a schema_path containing
    # `..` cannot escape the schemas directory downstream.
    base = env_cfg.dictionary_schema_path.strip().strip("/").rsplit("/", 1)[-1]
    stem, dot, ext = base.rpartition(".")
    return f"{stem}_{v}.{ext}" if dot else f"{base}_{v}"


def dictionary_local_relpath(env_cfg: EnvConfig, version: Optional[str] = None) -> str:
    """Repo-relative path that a pulled dictionary lands at."""
    return f"{DICT_SCHEMA_DIR_RELPATH}/{dictionary_filename(env_cfg, version)}"


@dataclass(frozen=True)
class StudyConfig:
    """Resolved settings for one study (in a given environment)."""

    key: str
    project_id: str
    program_id: str
    s3_metadata_path: str


def resolve_study(cfg: dict, study: str, env: str) -> StudyConfig:
    """Resolve a study against an environment.

    Tries ``<study>_<env_base>`` first (so ``ausdiab`` + ``staging`` →
    ``ausdiab_staging``), then the literal ``study`` key for back-compat.
    """
    studies = cfg.get("studies", {})
    candidates = [f"{study}_{env_base(env)}", study]
    for key in candidates:
        sc = studies.get(key)
        if sc:
            return StudyConfig(
                key=key,
                project_id=sc["project_id"],
                program_id=sc["program_id"],
                s3_metadata_path=sc["s3_metadata_path"],
            )
    raise ConfigError(
        f"Study '{study}' (env '{env}') not found. Tried {candidates}. "
        f"Valid studies: {', '.join(list_studies(cfg))}"
    )
