#!/usr/bin/env python3
"""Check that stubs/ match the real Open WebUI signatures.

Fetches the relevant Open WebUI source files from GitHub at a given ref
(default: the tag from `required_open_webui_version` in comfier.py), parses
them with `ast` (nothing is imported or installed) and compares function
signatures with the .pyi stubs: async-ness, parameter names/kinds and whether
each parameter has a default. Annotations are not compared.
"""

import argparse
import ast
import pathlib
import re
import sys
import urllib.request

ROOT = pathlib.Path(__file__).parent.parent
RAW = "https://raw.githubusercontent.com/open-webui/open-webui/{ref}/backend/open_webui/{path}.py"

# (stub file, upstream module path, owner class or None, function names)
CHECKS: list[tuple[str, str, str | None, list[str]]] = [
    ("models/config", "models/config", "Config", ["get"]),
    ("models/users", "models/users", "UsersTable", ["get_user_by_id"]),
    ("routers/images", "routers/images", None, ["get_image_config", "upload_image"]),
]

Fn = ast.FunctionDef | ast.AsyncFunctionDef


def default_ref() -> str:
    header = (ROOT / "comfier.py").read_text(encoding="utf-8")
    match = re.search(r"^required_open_webui_version:\s*(\S+)", header, re.MULTILINE)
    if not match:
        sys.exit("Could not find required_open_webui_version in comfier.py")
    return "v" + match.group(1).lstrip("v")


def fetch(ref: str, path: str) -> str:
    url = RAW.format(ref=ref, path=path)
    with urllib.request.urlopen(url) as resp:
        return resp.read().decode("utf-8")


def find_fn(tree: ast.Module, owner: str | None, name: str) -> Fn | None:
    body: list[ast.stmt] = tree.body
    if owner is not None:
        cls = next(
            (n for n in body if isinstance(n, ast.ClassDef) and n.name == owner), None
        )
        if cls is None:
            return None
        body = cls.body
    return next(
        (
            n
            for n in body
            if isinstance(n, ast.FunctionDef | ast.AsyncFunctionDef) and n.name == name
        ),
        None,
    )


def signature(fn: Fn, drop_self: bool) -> tuple[bool, list[tuple[str, str, bool]]]:
    """(is_async, [(name, kind, has_default)])"""
    a = fn.args
    params: list[tuple[str, str, bool]] = []
    positional = [*a.posonlyargs, *a.args]
    n_defaults = len(a.defaults)
    for i, arg in enumerate(positional):
        has_default = i >= len(positional) - n_defaults
        kind = "posonly" if arg in a.posonlyargs else "pos"
        params.append((arg.arg, kind, has_default))
    if a.vararg:
        params.append((a.vararg.arg, "vararg", False))
    for arg, default in zip(a.kwonlyargs, a.kw_defaults, strict=True):
        params.append((arg.arg, "kwonly", default is not None))
    if a.kwarg:
        params.append((a.kwarg.arg, "kwarg", False))
    if drop_self and params and params[0][0] in ("self", "cls"):
        params = params[1:]
    return isinstance(fn, ast.AsyncFunctionDef), params


def check_users_singleton(upstream: ast.Module, stub: ast.Module) -> list[str]:
    """`Users` must be an instance of UsersTable in both."""
    problems: list[str] = []
    real = any(
        isinstance(n, ast.Assign)
        and any(isinstance(t, ast.Name) and t.id == "Users" for t in n.targets)
        and isinstance(n.value, ast.Call)
        and isinstance(n.value.func, ast.Name)
        and n.value.func.id == "UsersTable"
        for n in upstream.body
    )
    stubbed = any(
        isinstance(n, ast.AnnAssign)
        and isinstance(n.target, ast.Name)
        and n.target.id == "Users"
        and isinstance(n.annotation, ast.Name)
        and n.annotation.id == "UsersTable"
        for n in stub.body
    )
    if not real:
        problems.append("models/users: upstream no longer has `Users = UsersTable()`")
    if not stubbed:
        problems.append("models/users: stub lacks `Users: UsersTable`")
    return problems


def is_static(fn: Fn) -> bool:
    return any(
        isinstance(d, ast.Name) and d.id == "staticmethod" for d in fn.decorator_list
    )


def main() -> int:
    parser = argparse.ArgumentParser(description=(__doc__ or "").splitlines()[0])
    parser.add_argument(
        "--ref", default=None, help="git ref/tag (default: from comfier.py)"
    )
    args = parser.parse_args()
    ref = args.ref or default_ref()
    print(f"Checking stubs against open-webui @ {ref}")

    problems: list[str] = []
    for stub_path, upstream_path, owner, names in CHECKS:
        stub_file = ROOT / "stubs" / "open_webui" / f"{stub_path}.pyi"
        stub_tree = ast.parse(stub_file.read_text(encoding="utf-8"))
        try:
            upstream_tree = ast.parse(fetch(ref, upstream_path))
        except Exception as err:
            problems.append(f"{upstream_path}: could not fetch/parse upstream ({err})")
            continue

        for name in names:
            label = f"{upstream_path}:{owner + '.' if owner else ''}{name}"
            real_fn = find_fn(upstream_tree, owner, name)
            stub_fn = find_fn(stub_tree, owner, name)
            if real_fn is None:
                problems.append(f"{label}: missing upstream")
                continue
            if stub_fn is None:
                problems.append(f"{label}: missing in stub")
                continue
            real_sig = signature(
                real_fn, drop_self=owner is not None and not is_static(real_fn)
            )
            stub_sig = signature(
                stub_fn, drop_self=owner is not None and not is_static(stub_fn)
            )
            if real_sig != stub_sig:
                problems.append(
                    f"{label}:\n    upstream: {real_sig}\n    stub:     {stub_sig}"
                )
            else:
                print(f"  ok  {label}")

        if owner == "UsersTable":
            problems += check_users_singleton(upstream_tree, stub_tree)

    if problems:
        print("\nMismatches:")
        for p in problems:
            print(f"  - {p}")
        return 1
    print("All stubs match.")
    return 0


if __name__ == "__main__":
    sys.exit(main())