from __future__ import annotations

import inspect
import string
import sys

# ruff: noqa: N806
from collections import OrderedDict
from inspect import isfunction, ismethoddescriptor
from pathlib import Path
from types import MethodType, ModuleType
from typing import TYPE_CHECKING, Any

import polars as pl

import narwhals as nw

if TYPE_CHECKING:
    from collections.abc import Iterator

LOWERCASE = tuple(string.ascii_lowercase)

if sys.version_info >= (3, 13):

    def _is_public_method_or_property(obj: Any) -> bool:
        return (
            isfunction(obj)
            or (isinstance(obj, (MethodType, property)) or ismethoddescriptor(obj))
        ) and obj.__name__.startswith(LOWERCASE)
else:

    def _is_public_method_or_property(obj: Any) -> bool:
        return (
            (isfunction(obj) or (isinstance(obj, MethodType) or ismethoddescriptor(obj)))
            and obj.__name__.startswith(LOWERCASE)
        ) or (isinstance(obj, property) and obj.fget.__name__.startswith(LOWERCASE))


def iter_api_reference_names(tp: type[Any]) -> Iterator[str]:
    for name, _ in inspect.getmembers(tp, _is_public_method_or_property):
        yield name


def iter_api_reference_names_dtypes(module: ModuleType) -> Iterator[str]:
    base = module.DType
    # NOTE: Special case, only non-dtype
    field_dtype = module.Field
    for name, _ in inspect.getmembers(
        module,
        lambda x: isinstance(x, type) and (issubclass(x, base) or x is field_dtype),
    ):
        yield name


def read_documented_members(source: str | Path) -> list[str]:
    MEMBERS_START = "members:\n"
    MEMBERS_PREFIX = "        - "
    DUNDER_PREFIX = "__"

    members: list[str] = []
    in_members_section = False
    with Path(source).open(encoding="utf-8") as fd:
        for line in fd:
            if in_members_section:
                if line.startswith(MEMBERS_PREFIX):
                    members.append(line.removeprefix(MEMBERS_PREFIX).strip())
                else:
                    in_members_section = False
            elif line.endswith(MEMBERS_START):
                in_members_section = True
    return [m for m in members if not m.startswith(DUNDER_PREFIX)]


ret = 0

NAMESPACES = {"dt", "str", "cat", "name", "list", "struct"}
EXPR_ONLY_METHODS = {"over", "map_batches"}
SERIES_ONLY_METHODS = {
    "arg_max",
    "arg_min",
    "arg_true",
    "dtype",
    "from_iterable",
    "from_numpy",
    "gather_every",
    "implementation",
    "is_empty",
    "is_sorted",
    "head",
    "hist",
    "item",
    "name",
    "rename",
    "sample",
    "scatter",
    "shape",
    "sort",
    "tail",
    "to_arrow",
    "to_dummies",
    "to_frame",
    "to_list",
    "to_native",
    "to_numpy",
    "to_pandas",
    "to_polars",
    "value_counts",
    "zip_with",
    "__iter__",
    "__contains__",
}
DIR_API_REF = Path("docs/api-reference")

files = {fp.stem for fp in Path("narwhals").iterdir()}

# Top level functions
top_level_functions = [
    i for i in nw.__all__ if not i[0].isupper() and i[0] != "_" and i not in files
]

documented = read_documented_members(DIR_API_REF / "narwhals.md")
if missing := set(top_level_functions).difference(documented).difference({"annotations"}):
    print("top-level functions: not documented")
    print(missing)
    ret = 1
if extra := set(documented).difference(top_level_functions):
    print("top-level functions: outdated")
    print(extra)
    ret = 1

# DataFrame methods
dataframe_methods = list(iter_api_reference_names(nw.DataFrame))
documented = read_documented_members(DIR_API_REF / "dataframe.md")
if missing := set(dataframe_methods).difference(documented):
    print("DataFrame: not documented")
    print(missing)
    ret = 1
if extra := set(documented).difference(dataframe_methods):
    print("DataFrame: outdated")
    print(extra)
    ret = 1

# LazyFrame methods
lazyframe_methods = list(iter_api_reference_names(nw.LazyFrame))
documented = read_documented_members(DIR_API_REF / "lazyframe.md")
if missing := set(lazyframe_methods).difference(documented):
    print("LazyFrame: not documented")
    print(missing)
    ret = 1
if extra := set(documented).difference(lazyframe_methods):
    print("LazyFrame: outdated")
    print(extra)
    ret = 1

# Series methods
series_methods = list(iter_api_reference_names(nw.Series))
documented = read_documented_members(DIR_API_REF / "series.md")
if missing := set(series_methods).difference(documented).difference(NAMESPACES):
    print("Series: not documented")
    print(missing)
    ret = 1
if extra := set(documented).difference(series_methods):
    print("Series: outdated")
    print(extra)
    ret = 1

# Series.{cat, dt, list, str} methods
for namespace in NAMESPACES.difference({"name"}):
    series_ns_methods = [
        i
        for i in dir(getattr(nw.from_native(pl.Series(), series_only=True), namespace))
        if not i[0].isupper() and i[0] != "_"
    ]
    documented = read_documented_members(DIR_API_REF / f"series_{namespace}.md")
    if missing := set(series_ns_methods).difference(documented):
        print(f"Series.{namespace}: not documented")
        print(missing)
        ret = 1
    if extra := set(documented).difference(series_ns_methods):
        print(f"Series.{namespace}: outdated")
        print(extra)
        ret = 1

# Expr methods
expr_methods = list(iter_api_reference_names(nw.Expr))
documented = read_documented_members(DIR_API_REF / "expr.md")
if missing := set(expr_methods).difference(documented).difference(NAMESPACES):
    print("Expr: not documented")
    print(missing)
    ret = 1
if extra := set(documented).difference(expr_methods):
    print("Expr: outdated")
    print(extra)
    ret = 1

# Expr.{cat, dt, list, name, str} methods
for namespace in NAMESPACES:
    expr_ns_methods = [
        i
        for i in dir(getattr(nw.col("a"), namespace))
        if not i[0].isupper() and i[0] != "_"
    ]
    documented = read_documented_members(DIR_API_REF / f"expr_{namespace}.md")
    if missing := set(expr_ns_methods).difference(documented):
        print(f"Expr.{namespace}: not documented")
        print(missing)
        ret = 1
    if extra := set(documented).difference(expr_ns_methods):
        print(f"Expr.{namespace}: outdated")
        print(extra)
        ret = 1

# DTypes
dtypes = list(iter_api_reference_names_dtypes(nw.dtypes))
documented = read_documented_members(DIR_API_REF / "dtypes.md")
if missing := set(dtypes).difference(documented):
    print("Dtype: not documented")
    print(missing)
    ret = 1
if extra := set(documented).difference(dtypes):
    print("Dtype: outdated")
    print(extra)
    ret = 1

# Schema
schema_methods = list(iter_api_reference_names(nw.Schema))
documented = read_documented_members(DIR_API_REF / "schema.md")
if (
    missing := set(schema_methods)
    .difference(documented)
    .difference(iter_api_reference_names(OrderedDict))
):
    print("Schema: not documented")
    print(missing)
    ret = 1
if extra := set(documented).difference(schema_methods):
    print("Schema: outdated")
    print(extra)
    ret = 1

# Check Expr vs Series
if missing := set(expr_methods).difference(series_methods).difference(EXPR_ONLY_METHODS):
    print("In Expr but not in Series")
    print(missing)
    ret = 1
if extra := set(series_methods).difference(expr_methods).difference(SERIES_ONLY_METHODS):
    print("In Series but not in Expr")
    print(extra)
    ret = 1

# Check Expr vs Series internal methods
for namespace in NAMESPACES.difference({"name"}):
    expr_internal = [
        i
        for i in dir(getattr(nw.col("a"), namespace))
        if not i[0].isupper() and i[0] != "_"
    ]
    series_internal = [
        i
        for i in dir(getattr(nw.from_native(pl.Series(), series_only=True), namespace))
        if not i[0].isupper() and i[0] != "_"
    ]
    if missing := set(expr_internal).difference(series_internal):
        print(f"In Expr.{namespace} but not in Series.{namespace}")
        print(missing)
        ret = 1
    if extra := set(series_internal).difference(expr_internal):
        print(f"In Series.{namespace} but not in Expr.{namespace}")
        print(extra)
        ret = 1

sys.exit(ret)
