from __future__ import annotations


def generate_tabular_overloads(func_name: str) -> str:
    """Generates overloads for a function with the specified engines and lazy options."""
    overloads = []
    ENGINES = ["pandas", "polars", "dask"]
    LAZY_OPTIONS = [False, True]

    for engine in ENGINES:
        for lazy in LAZY_OPTIONS:
            if engine == "pandas" and not lazy:
                return_type = "_pd.DataFrame"
            elif engine == "polars":
                return_type = "_pl.LazyFrame" if lazy else "_pl.DataFrame"
            elif engine == "dask" and lazy:
                return_type = "_dd.DataFrame"
            else:
                continue

            # Build the overload definition using .format
            overload_def = """\
@overload
def {func_name}(
    engine: {engine_literal},
    *,
    engine_kwargs: dict[str, Any] | None = None,
    lazy: {lazy_literal} = {lazy_repr},
) -> {return_type}: ...""".format(
                func_name=func_name,
                engine_literal=f'Literal["{engine}"]' if engine else "None",
                lazy_literal=f"Literal[{lazy}]",
                lazy_repr=lazy,
                return_type=return_type,
            )
            overloads.append(overload_def)

    return "\n".join(overloads)


def generate_gridded_overloads(func_name: str) -> str:
    """Generates overloads for a function with the specified engines and lazy options."""
    overloads = []
    ENGINES = ["xarray"]
    LAZY_OPTIONS = [False]

    for engine in ENGINES:
        for lazy in LAZY_OPTIONS:
            if engine is None:
                return_type = "" if lazy else "_xr.Dataset"
            elif engine == "xarray":
                return_type = "_xr.Dataset"
            else:
                continue

            # Build the overload definition using .format
            overload_def = """\
@overload
def {func_name}(
    engine: {engine_literal},
    *,
    engine_kwargs: dict[str, Any] | None = None,
    # lazy: {lazy_literal} = {lazy_repr},
) -> {return_type}: ...""".format(
                func_name=func_name,
                engine_literal=f'Literal["{engine}"]' if engine else "None",
                lazy_literal=f"Literal[{lazy}]",
                lazy_repr=lazy,
                return_type=return_type,
            )
            overloads.append(overload_def)

    return "\n".join(overloads)


synthetic_clusters = """@overload
def synthetic_clusters(
    engine: Literal["pandas"],
    *,
    lazy: Literal[False] = False,
    total_points: int = 1_000_000,
) -> _pd.DataFrame: ...
@overload
def synthetic_clusters(
    engine: Literal["polars"],
    *,
    lazy: Literal[False] = False,
    total_points: int = 1_000_000,
) -> _pl.DataFrame: ...
@overload
def synthetic_clusters(
    engine: Literal["polars"],
    *,
    lazy: Literal[True] = True,
    total_points: int = 1_000_000,
) -> _pl.LazyFrame: ...
@overload
def synthetic_clusters(
    engine: Literal["dask"],
    *,
    lazy: bool = False,
    total_points: int = 1_000_000,
) -> _dd.DataFrame: ...
"""
landsat_rgb = """@overload
def landsat_rgb(
    engine: Literal["rioxarray"],
    *,
    engine_kwargs: dict[str, Any] | None = None,
) -> _xr.Dataset: ...
"""
penguins_rgba = """@overload
def penguins_rgba(
    engine: Literal["xarray"],
) -> _xr.Dataset: ...
"""
us_states = """@overload
def us_states(
    engine: Literal["geopandas"],
    *,
    engine_kwargs: dict[str, Any] | None = None,
) -> _gpd.GeoDataFrame: ...
"""


def custom() -> str:
    return "\n".join(  # noqa: FLY002
        [
            synthetic_clusters.strip(),
            landsat_rgb.strip(),
            us_states.strip(),
        ]
    )


def main():
    header = """\
# !!! THIS FILE IS AUTOGENERATED !!!
# For changes see scripts/generate_typehints.py

from __future__ import annotations

from typing import Any, Literal, overload

import dask.dataframe as _dd
import geopandas as _gpd
import pandas as _pd
import polars as _pl
import xarray as _xr

__version__: str
__all__: tuple[str, ...]
"""

    return "\n".join(
        [
            header,
            generate_tabular_overloads("penguins"),
            generate_tabular_overloads("earthquakes"),
            generate_tabular_overloads("apple_stocks"),
            generate_tabular_overloads("stocks"),
            generate_tabular_overloads("nyc_taxi_remote"),
            generate_gridded_overloads("air_temperature"),
            custom(),
            "",
        ]
    )


if __name__ == "__main__":
    with open("src/hvsampledata/__init__.pyi", "w") as f:
        f.write(main())
