from pathlib import Path
from os import PathLike
from typing import Any, Literal, Sequence, overload
import yaml as yaml_lib
import tomllib
import json as json_lib
from pydantic import BaseModel

DEFAULT_IMPORT_KEY = "import"
DEFAULT_OVERRIDE_KEY = "override"

type FileTypes = list[Literal["yaml", "toml", "json"]] | Literal["all"]


def __load_raw(path: PathLike, file_types: FileTypes) -> dict[str, Any]:
    supported_exts = []
    if file_types == "all":
        supported_exts = [".yaml", ".yml", ".toml", ".json"]
    else:
        if "yaml" in file_types:
            supported_exts.extend([".yaml", ".yml"])
        if "toml" in file_types:
            supported_exts.append(".toml")
        if "json" in file_types:
            supported_exts.append(".json")

    path = Path(path)
    if path.suffix.lower() not in supported_exts:
        raise ValueError(f"Unsupported file format: {path}")
    if path.suffix.lower() in [".yaml", ".yml"]:
        with open(path, "r") as f:
            data = yaml_lib.safe_load(f)
    elif path.suffix.lower() == ".toml":
        with open(path, "r") as f:
            s = f.read()
            data = tomllib.loads(s)
    elif path.suffix.lower() == ".json":
        with open(path, "r") as f:
            data = json_lib.load(f)
    else:
        raise ValueError(f"Unsupported file format: {path}")
    if not isinstance(data, dict):
        raise ValueError(f"Invalid config {path}: root element must be a dictionary.")
    return data


def __merge_docs(doc1: dict[str, Any], doc2: dict[str, Any]) -> dict[str, Any]:
    """
    Merges two dictionaries, with dict2 taking precedence over dict1.
    """
    result = doc1.copy()
    for key, value in doc2.items():
        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
            result[key] = __merge_docs(result[key], value)
        else:
            result[key] = value
    return result


def __apply_overrides(data: dict[str, Any], overrides: dict[str, Any]):
    for key, value in overrides.items():
        keys = key.split(".")
        d = data
        for k in keys[:-1]:
            if k not in d or not isinstance(d, dict):
                raise KeyError(f"Key {key} not found in configuration.")
            d = d[k]
        if not isinstance(d, dict):
            raise KeyError(f"Key {key} not found in configuration.")
        d[keys[-1]] = value
    return data


def __resolve_doc(
    path: PathLike, import_key: str, override_key: str, file_types: FileTypes
) -> Any:
    data = __load_raw(path, file_types)
    # Recursively resolve imports
    if import_key in data:
        paths = (
            data[import_key]
            if isinstance(data[import_key], list)
            else [data[import_key]]
        )
        for p in paths:
            if not isinstance(p, str):
                raise ValueError(f"Invalid import path: {p}")
            import_path = Path(p)
            if not import_path.is_absolute():
                import_path = Path(path).parent / p
            imported_data = __resolve_doc(
                import_path, import_key, override_key, file_types
            )
            data = __merge_docs(imported_data, data)
        del data[import_key]
    # Apply overrides
    if override_key in data:
        overrides = data[override_key]
        __apply_overrides(data, overrides)
        del data[override_key]
    return data


type RawObject = dict[str, "RawObject"] | list[
    "RawObject"
] | str | int | float | bool | None

type RawDoc = dict[str, RawObject]


@overload
def load[M: BaseModel](
    path: PathLike,
    *,
    type: type[M],
    import_key: str = DEFAULT_IMPORT_KEY,
    override_key: str = DEFAULT_OVERRIDE_KEY,
    file_types: FileTypes = "all",
) -> M: ...


@overload
def load[M: BaseModel](
    path: PathLike,
    *,
    type: None = None,
    import_key: str = DEFAULT_IMPORT_KEY,
    override_key: str = DEFAULT_OVERRIDE_KEY,
    file_types: FileTypes = "all",
) -> RawDoc: ...


def load[M: BaseModel](
    path: PathLike,
    *,
    type: type[M] | None = None,
    import_key: str = DEFAULT_IMPORT_KEY,
    override_key: str = DEFAULT_OVERRIDE_KEY,
    file_types: FileTypes = "all",
) -> RawDoc | M:
    doc = __resolve_doc(path, import_key, override_key, file_types)
    if type is not None:
        return type.model_validate(doc)
    return doc


@overload
def yaml[M: BaseModel](
    path: PathLike,
    *,
    type: type[M],
    import_key: str = DEFAULT_IMPORT_KEY,
    override_key: str = DEFAULT_OVERRIDE_KEY,
) -> M: ...


@overload
def yaml(
    path: PathLike,
    *,
    type: None = None,
    import_key: str = DEFAULT_IMPORT_KEY,
    override_key: str = DEFAULT_OVERRIDE_KEY,
) -> RawDoc: ...


def yaml[M: BaseModel](
    path: PathLike,
    *,
    type: type[M] | None = None,
    import_key: str = DEFAULT_IMPORT_KEY,
    override_key: str = DEFAULT_OVERRIDE_KEY,
) -> RawDoc | M:
    doc = __resolve_doc(path, import_key, override_key, file_types=["yaml"])
    if type is not None:
        return type.model_validate(doc)
    return doc


@overload
def toml[M: BaseModel](
    path: PathLike,
    *,
    type: type[M],
    import_key: str = DEFAULT_IMPORT_KEY,
    override_key: str = DEFAULT_OVERRIDE_KEY,
) -> M: ...


@overload
def toml(
    path: PathLike,
    *,
    type: None = None,
    import_key: str = DEFAULT_IMPORT_KEY,
    override_key: str = DEFAULT_OVERRIDE_KEY,
) -> RawDoc: ...


def toml[M: BaseModel](
    path: PathLike,
    *,
    type: type[M] | None = None,
    import_key: str = DEFAULT_IMPORT_KEY,
    override_key: str = DEFAULT_OVERRIDE_KEY,
) -> RawDoc | M:
    doc = __resolve_doc(path, import_key, override_key, file_types=["toml"])
    if type is not None:
        return type.model_validate(doc)
    return doc


@overload
def json[M: BaseModel](
    path: PathLike,
    *,
    type: type[M],
    import_key: str = DEFAULT_IMPORT_KEY,
    override_key: str = DEFAULT_OVERRIDE_KEY,
) -> M: ...


@overload
def json(
    path: PathLike,
    *,
    type: None = None,
    import_key: str = DEFAULT_IMPORT_KEY,
    override_key: str = DEFAULT_OVERRIDE_KEY,
) -> RawDoc: ...


def json[M: BaseModel](
    path: PathLike,
    *,
    type: type[M] | None = None,
    import_key: str = DEFAULT_IMPORT_KEY,
    override_key: str = DEFAULT_OVERRIDE_KEY,
) -> RawDoc | M:
    doc = __resolve_doc(path, import_key, override_key, file_types=["json"])
    if type is not None:
        return type.model_validate(doc)
    return doc
