"""Module for reading yera config toml."""

from pathlib import Path

import tomllib

from yera.config2.dataclasses import YeraConfig


def _read_toml(use_global: bool = False) -> dict:
    """Read yera.toml configuration file.

    First this function tries in CWD, then HOME/.config/yera and then errors if neither
    are present.

    Args:
        use_global: whether to use the global config, if local is present (default:
        False).

    Returns:
        dict: yera.toml content loaded into a dictionary

    """
    toml_path = Path.cwd() / "yera.toml"
    if not toml_path.exists() or use_global:
        toml_path = Path.home() / ".config" / "yera" / "yera.toml"

    if not toml_path.exists():
        raise FileNotFoundError(
            "No yera.toml found in current working directory or home/.config/yera"
        )

    with toml_path.open("rb") as f:
        return tomllib.load(f)


def _fill_config_object(config_dict: dict) -> YeraConfig:
    """Create YeraConfig from parsed TOML dictionary.

    This method flattens the nested TOML structure (models.llm.provider.model_name)
    into the flat dictionary structure (models.llm["provider.model_name"]) and
    adds the `id` field to each model.

    Args:
        config_dict: Dictionary from tomllib.load()

    Returns:
        Validated YeraConfig instance

    """

    def _flatten_nested_models(
        nested_dict: dict, path: list[str] | None = None
    ) -> dict:
        if path is None:
            path = []

        flattened = {}

        for k, v in nested_dict.items():
            current_path = [*path, k]

            is_leaf = isinstance(v, dict) and any(
                not isinstance(sv, dict) for sv in v.values()
            )

            if is_leaf:
                full_id = ".".join(current_path)
                model_cfg = {**v, "id": full_id}
                flattened[full_id] = model_cfg
            elif isinstance(v, dict):
                flattened.update(_flatten_nested_models(v, current_path))
            else:
                raise TypeError(
                    f"Expected nested dicts only in models section. Encountered {type(v)}"
                )
        return flattened

    kws = {**config_dict, "models": {}}

    for model_type in config_dict["models"]:
        kws["models"][model_type] = _flatten_nested_models(
            config_dict["models"][model_type]
        )

    return YeraConfig(**kws)


def read_config(use_global: bool = False) -> YeraConfig:
    """Read and parse yera.toml configuration into a YeraConfig object.

    Searches for yera.toml in the current working directory first, then falls
    back to ~/.config/yera/yera.toml. The TOML structure is parsed and validated
    into a YeraConfig dataclass instance.

    Args:
        use_global: If True, skip the local directory and use the global config
            at ~/.config/yera/yera.toml even if a local config exists (default: False).

    Returns:
        Validated YeraConfig instance with flattened model configurations.

    Raises:
        FileNotFoundError: If no yera.toml is found in either location.
        ValidationError: If the TOML structure doesn't match YeraConfig schema.

    Example:
        >>> config = read_config()  # Uses local config if available
        >>> config = read_config(use_global=True)  # Forces global config

    """
    toml_dict = _read_toml(use_global=use_global)
    return _fill_config_object(toml_dict)
