from pathlib import Path

from huggingface_hub import hf_hub_download
from omegaconf import DictConfig
from omegaconf import OmegaConf
from torch import Tensor

from .defaults import REPO_NAME
from .defaults import REPO_USER
from .defaults import REVISION
from .defaults import ROOT_CONFIG_FILENAME
from .defaults import WEIGHTS_FILENAME
from .defaults import get_configs_dir

TypeStateDict = dict[str, Tensor]


def download_weights(**kwargs) -> Path:
    weights_path = _download_from_hugging_face(**kwargs)
    return weights_path


def _download_from_hugging_face(
    repo_user: str = REPO_USER,
    repo_name: str = REPO_NAME,
    revision: str | None = REVISION,
    filename: str = WEIGHTS_FILENAME,
) -> Path:
    repo_id = f"{repo_user}/{repo_name}"
    try:
        weights_path = hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            revision=revision,
        )
    except Exception as e:
        msg = f'Failed to download "{filename}" from Hugging Face Hub repo "{repo_id}".'
        raise RuntimeError(msg) from e
    return Path(weights_path)


def _resolve_defaults(config_dir: Path, config_name: str) -> DictConfig:
    """Load a YAML config and recursively resolve its Hydra-style defaults."""
    config_path = config_dir / f"{config_name}.yaml"
    raw = OmegaConf.load(config_path)
    assert isinstance(raw, DictConfig)

    defaults = OmegaConf.to_container(raw.pop("defaults", OmegaConf.create([])))
    assert isinstance(defaults, list)

    merged = OmegaConf.create()
    for default in defaults:
        if default == "_self_":
            continue
        if isinstance(default, dict):
            for group, option in default.items():
                sub = _resolve_defaults(config_dir / group, str(option))
                merged = OmegaConf.merge(merged, OmegaConf.create({group: sub}))

    # _self_ is always applied last (standard Hydra behavior)
    merged = OmegaConf.merge(merged, raw)
    assert isinstance(merged, DictConfig)
    return merged


def _load_config(*, overrides: list[str] | None = None) -> DictConfig:
    config_name = ROOT_CONFIG_FILENAME.removesuffix(".yaml")
    config = _resolve_defaults(get_configs_dir(), config_name)
    if overrides:
        config = OmegaConf.merge(config, OmegaConf.from_dotlist(overrides))
    assert isinstance(config, DictConfig)
    return config


def load_model_config(**kwargs) -> DictConfig:
    return _load_config(**kwargs)["model"]


def load_processor_config(**kwargs) -> DictConfig:
    return _load_config(**kwargs)["processor"]
