"""Module to interact with the AHORN dataset API."""

import contextlib
import gzip
import json
import logging
from collections.abc import AsyncIterator, Iterator
from datetime import UTC, datetime
from pathlib import Path
from typing import TypedDict

import httpx
from httpx_retries import Retry, RetryTransport

from .utils import get_cache_dir
from .validator import Validator

__all__ = [
    "download_dataset",
    "get_dataset_url",
    "load_dataset_data",
    "load_datasets_data",
    "read_dataset",
    "validate_dataset",
]

DATASET_API_URL = "https://ahorn.rwth-aachen.de/api/datasets.json"

logger = logging.getLogger(__name__)


class AttachmentDict(TypedDict):
    url: str
    size: int


class DatasetDict(TypedDict):
    slug: str
    title: str
    tags: list[str]
    attachments: dict[str, AttachmentDict]


class DatasetsDataDict(TypedDict):
    datasets: dict[str, DatasetDict]
    time: str


async def load_datasets_data(
    *, cache_lifetime: int | None = None
) -> dict[str, DatasetDict]:
    """Load dataset data from the Ahorn API.

    Parameters
    ----------
    cache_lifetime : int, optional
        How long to reuse cached data in seconds. If not provided, the cache will not
        be used.

    Returns
    -------
    dict[str, Any]
        Dictionary containing dataset information, where the keys are dataset slugs
        and the values are dictionaries with dataset details such as title, tags, and
        attachments.
    """
    datasets_data_cache = get_cache_dir() / "datasets.json"
    if datasets_data_cache.exists() and cache_lifetime is not None:
        cache_mtime = datetime.fromtimestamp(
            datasets_data_cache.stat().st_mtime, tz=UTC
        )
        age_seconds = (datetime.now(tz=UTC) - cache_mtime).total_seconds()
        if age_seconds < cache_lifetime:
            logger.info(
                "Using cached datasets list (age=%.1fs, lifetime=%ss)",
                age_seconds,
                cache_lifetime,
            )
            with datasets_data_cache.open("r", encoding="utf-8") as cache_file:
                cache: DatasetsDataDict = json.load(cache_file)
                return cache["datasets"]

    logger.info("Fetching datasets list from %s", DATASET_API_URL)
    async with httpx.AsyncClient(timeout=10) as client:
        response = await client.get(DATASET_API_URL)
        response.raise_for_status()

    datasets_data_cache.parent.mkdir(parents=True, exist_ok=True)
    with datasets_data_cache.open("w", encoding="utf-8") as cache_file:
        cache_file.write(response.text)

    response_json: DatasetsDataDict = response.json()
    return response_json["datasets"]


async def load_dataset_data(
    slug: str, *, cache_lifetime: int | None = None
) -> DatasetDict:
    """Load data for a specific dataset by its slug.

    Parameters
    ----------
    slug : str
        The slug of the dataset to load.
    cache_lifetime : int, optional
        How long to reuse cached data in seconds. If not provided, the cache will not
        be used.

    Returns
    -------
    DatasetDict
        Dictionary containing the dataset details.

    Raises
    ------
    ValueError
        If the dataset with the given `slug` does not exist.
    """
    datasets = await load_datasets_data(cache_lifetime=cache_lifetime)

    if slug not in datasets:
        raise ValueError(f"Dataset with slug '{slug}' does not exist in AHORN.")

    return datasets[slug]


async def get_dataset_url(
    slug: str, revision: int | None = None, *, cache_lifetime: int | None = None
) -> httpx.URL:
    """Get the download URL for a specific dataset by its slug.

    Parameters
    ----------
    slug : str
        The slug of the dataset.
    revision : int, optional
        The revision number to download. If not provided, the latest revision will be
        used.
    cache_lifetime : int, optional
        How long to reuse cached data in seconds. If not provided, the cache will not
        be used.

    Returns
    -------
    httpx.URL
        The download URL of the dataset.

    Raises
    ------
    ValueError
        If the dataset with the given `slug` does not exist.
    ValueError
        If the revision number is not valid, e.g. because the specified revision does
        not exist for the dataset.
    RuntimeError
        If the dataset does not contain the required attachment information.

    Examples
    --------
    >>> import asyncio
    >>> asyncio.run(get_dataset_url("karate-club", revision=2))
    URL('https://zenodo.org/records/19236080/files/karate-club.txt?download=1')
    """
    data = await load_dataset_data(slug, cache_lifetime=cache_lifetime)

    revisions = [key for key in data["attachments"] if key.startswith("revision-")]

    if not revisions:
        raise RuntimeError(
            f"Dataset '{slug}' does not contain any revision attachments. "
            "This is an error with the dataset metadata and should be reported to the maintainers."
        )

    if revision is None:
        # Use the latest revision (highest number)
        revision_numbers = [int(key.split("-")[1]) for key in revisions]
        latest_revision = max(revision_numbers)
        revision_key = f"revision-{latest_revision}"
        logger.info(
            "No revision was explicitly specified, using latest revision: %d",
            latest_revision,
        )
    else:
        revision_key = f"revision-{revision}"
        if revision_key not in data["attachments"]:
            available = sorted([int(k.split("-")[1]) for k in revisions])
            raise ValueError(
                f"Dataset '{slug}' does not have revision {revision}. "
                f"Available revisions: {available}"
            )

    return httpx.URL(data["attachments"][revision_key]["url"])


async def download_dataset(
    slug: str,
    folder: Path | str,
    revision: int | None = None,
    *,
    cache_lifetime: int | None = None,
) -> Path:
    """Download a dataset by its slug to the specified folder.

    This function implements an exponential backoff strategy when encountering HTTP 429
    (Too Many Requests) responses. If available, it respects the 'Retry-After' header to
    determine the wait time before retrying.

    Parameters
    ----------
    slug : str
        The slug of the dataset to download.
    folder : Path | str
        The folder where the dataset should be saved.
    revision : int, optional
        The revision number to download. If not provided, the latest revision will be
        used. Specifying an explicit revision can be important for reproducibility.
    cache_lifetime : int, optional
        How long to reuse cached data in seconds. If not provided, the cache will not
        be used.

    Returns
    -------
    Path
        The path to the downloaded dataset file.

    Raises
    ------
    ValueError
        If the dataset with the given `slug` does not exist.
    ValueError
        If the revision number is not valid, e.g. because the specified revision does
        not exist for the dataset.
    HTTPError
        If the dataset file could not be downloaded due to some error.
    RuntimeError
        If the dataset file could not be downloaded due to some error.
    """
    if isinstance(folder, str):
        folder = Path(folder)

    logger.info("Preparing download for dataset '%s' into %s", slug, folder)
    download_url = await get_dataset_url(slug, revision, cache_lifetime=cache_lifetime)

    folder.mkdir(parents=True, exist_ok=True)
    filepath = folder / download_url.path.split("/")[-1]

    # Use RetryTransport to automatically handle rate limiting (429) with exponential
    # backoff. This also automatically respects 'Retry-After' headers if provided.
    retry = Retry(
        total=5,
        backoff_factor=2.0,
    )
    retry_transport = RetryTransport(retry=retry)

    async with (
        httpx.AsyncClient(transport=retry_transport, timeout=10) as client,
        client.stream("GET", download_url) as response,
    ):
        response.raise_for_status()

        with filepath.open("wb") as f:
            async for chunk in response.aiter_bytes(chunk_size=8192):
                f.write(chunk)

    logger.info("Downloaded dataset '%s' to %s", slug, filepath)

    return filepath


@contextlib.asynccontextmanager
async def read_dataset(
    slug: str, revision: int | None = None
) -> AsyncIterator[Iterator[str]]:
    """Download and yield a context-managed file object for the dataset lines by slug.

    The dataset file will be stored in your system cache and can be deleted according
    to your system's cache policy. To ensure that costly re-downloads do not occur, use
    the `download_dataset` function to store the dataset file at a more permanent
    location.

    Parameters
    ----------
    slug : str
        The slug of the dataset to download.
    revision : int, optional
        The revision number to download. If not provided, the latest revision will be
        used. Specifying an explicit revision can be important for reproducibility.

    Returns
    -------
    Async context manager yielding an open file object (iterator over lines).

    Raises
    ------
    ValueError
        If the dataset with the given `slug` does not exist.
    ValueError
        If the revision number is not valid, e.g. because the specified revision does
        not exist for the dataset.
    RuntimeError
        If the dataset file could not be downloaded due to other errors.

    Examples
    --------
    >>> async def _read_first_line() -> str:
    ...     async with ahorn_loader.read_dataset("karate-club", revision=2) as dataset:
    ...         return next(dataset).strip()
    >>> import asyncio
    >>> asyncio.run(_read_first_line())
    '{"name": "karate-club", "format-version": "0.3", "revision": 2}'
    """
    download_url = await get_dataset_url(slug, revision)
    filepath = get_cache_dir() / download_url.path.split("/")[-1]

    # Download the dataset if it is not already cached
    if not filepath.exists():
        filepath = await download_dataset(slug, get_cache_dir(), revision)

    if filepath.suffix == ".gz":
        with gzip.open(filepath, mode="rt", encoding="utf-8") as f:
            yield f
    else:
        with filepath.open("r", encoding="utf-8") as f:
            yield f


def validate_dataset(path: str | Path) -> bool:
    r"""Validate whether a given file is a valid AHORN dataset.

    Validation errors are logged, but not raised as exceptions.

    Parameters
    ----------
    path : str or pathlib.Path
        The path to the dataset file to validate.

    Examples
    --------
    >>> from pathlib import Path
    >>> from tempfile import TemporaryDirectory
    >>> with TemporaryDirectory() as tmp_dir:
    ...     dataset_path = Path(tmp_dir) / "demo.txt"
    ...     _ = dataset_path.write_text(
    ...         '{"name": "demo", "revision": 1, "format-version": "0.1"}\n'
    ...         '1 {"weight": 1.0}\n'
    ...         '1,2 {"weight": 2.0}\n'
    ...     )
    ...     ahorn_loader.validate_dataset(dataset_path)
    True
    """
    validator = Validator()
    return validator.validate(path)
