# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Common utilities for the `tvm-ffi-stubgen` tool."""

from __future__ import annotations

import dataclasses
from io import StringIO
from typing import Any, Callable

from tvm_ffi.core import TypeInfo, TypeSchema, _lookup_type_attr

from . import consts as C


def _parse_type_schema(raw: str | dict[str, Any]) -> TypeSchema:
    """Parse a type schema from either a JSON string or an already-parsed dict."""
    if isinstance(raw, dict):
        return TypeSchema.from_json_obj(raw)
    return TypeSchema.from_json_str(raw)


@dataclasses.dataclass
class InitConfig:
    """Configuration for generating new stubs.

    Examples
    --------
    If we are generating type stubs for Python package `my-ffi-extension`,
    and the CMake target that generates the shared library is `my_ffi_extension_shared`,
    then we can run the following command to generate the stubs:

    --init-pypkg my-ffi-extension --init-lib my_ffi_extension_shared --init-prefix my_ffi_extension.

    """

    pkg: str
    """Name of the Python package to generate stubs for, e.g. apache-tvm-ffi (instead of tvm_ffi)"""

    shared_target: str
    """Name of CMake target that generates the shared library, e.g. tvm_ffi_shared

    This is used to determine the name of the shared library file.
    - macOS: lib{shared_target}.dylib or lib{shared_target}.so
    - Linux: lib{shared_target}.so
    - Windows: {shared_target}.dll
    """

    prefix: str
    """Only generate stubs for global function and objects with the given prefix, e.g. `tvm_ffi.`"""


@dataclasses.dataclass
class Options:
    """Command line options for stub generation."""

    imports: list[str] = dataclasses.field(default_factory=list)
    dlls: list[str] = dataclasses.field(default_factory=list)
    init: InitConfig | None = None
    indent: int = 4
    files: list[str] = dataclasses.field(default_factory=list)
    verbose: bool = False
    dry_run: bool = False


@dataclasses.dataclass(frozen=True, eq=True)
class ImportItem:
    """An import statement item."""

    mod: str
    name: str
    type_checking_only: bool = False
    alias: str | None = None

    def __init__(
        self,
        name: str,
        type_checking_only: bool = False,
        alias: str | None = None,
    ) -> None:
        """Initialize an `ImportItem` with the given module name and optional alias."""
        if "." in name:
            mod, name = name.rsplit(".", 1)
            for mod_prefix, mod_replacement in C.MOD_MAP.items():
                if mod.startswith(mod_prefix):
                    mod = mod.replace(mod_prefix, mod_replacement, 1)
                    break
        else:
            mod = ""
        object.__setattr__(self, "mod", mod)
        object.__setattr__(self, "name", name)
        object.__setattr__(self, "type_checking_only", type_checking_only)
        object.__setattr__(self, "alias", alias)

    @property
    def name_with_alias(self) -> str:
        """Generate a string of the form `name as alias` if an alias is set, otherwise just `name`."""
        return f"{self.name} as {self.alias}" if self.alias else self.name

    @property
    def full_name(self) -> str:
        """Generate a string of the form `mod.name` or `name` if no module is set."""
        return f"{self.mod}.{self.name}" if self.mod else self.name

    def __repr__(self) -> str:
        """Generate an import statement string for this item."""
        return str(self)

    def __str__(self) -> str:
        """Generate an import statement string for this item."""
        if self.mod:
            ret = f"from {self.mod} import {self.name_with_alias}"
        else:
            ret = f"import {self.name_with_alias}"
        return ret


@dataclasses.dataclass(init=False)
class NamedTypeSchema(TypeSchema):
    """A type schema with an associated name."""

    name: str

    def __init__(self, name: str, schema: TypeSchema) -> None:
        """Initialize a `NamedTypeSchema` with the given name and type schema."""
        super().__init__(origin=schema.origin, args=schema.args)
        self.name = name


@dataclasses.dataclass
class FuncInfo:
    """Information of a function."""

    schema: NamedTypeSchema
    is_member: bool

    @staticmethod
    def from_schema(name: str, schema: TypeSchema, *, is_member: bool = False) -> FuncInfo:
        """Construct a `FuncInfo` from a name and its type schema."""
        return FuncInfo(schema=NamedTypeSchema(name=name, schema=schema), is_member=is_member)

    def gen(self, ty_map: Callable[[str], str], indent: int) -> str:
        """Generate a function signature string for this function."""
        try:
            _, func_name = self.schema.name.rsplit(".", 1)
        except ValueError:
            func_name = self.schema.name
        buf = StringIO()
        buf.write(" " * indent)
        buf.write(f"def {func_name}(")
        if self.schema.origin != "Callable":
            raise ValueError(f"Expected Callable type schema, but got: {self.schema}")
        if not self.schema.args:
            ty_map("Any")
            buf.write("*args: Any) -> Any: ...")
            return buf.getvalue()
        arg_ret = self.schema.args[0]
        arg_args = self.schema.args[1:]
        for i, arg in enumerate(arg_args):
            if self.is_member and i == 0:
                buf.write("self, ")
            else:
                buf.write(f"_{i}: ")
                buf.write(arg.repr(ty_map))
                buf.write(", ")
        if arg_args:
            buf.write("/")
        buf.write(") -> ")
        buf.write(arg_ret.repr(ty_map))
        buf.write(": ...")
        return buf.getvalue()


@dataclasses.dataclass
class InitFieldInfo:
    """A field that participates in the auto-generated ``__init__``."""

    name: str
    schema: NamedTypeSchema
    kw_only: bool
    has_default: bool


@dataclasses.dataclass
class ObjectInfo:
    """Information of an object type, including its fields and methods."""

    fields: list[NamedTypeSchema]
    methods: list[FuncInfo]
    type_key: str | None = None
    parent_type_key: str | None = None
    init_fields: list[InitFieldInfo] = dataclasses.field(default_factory=list)
    has_init: bool = False

    @staticmethod
    def from_type_info(type_info: TypeInfo) -> ObjectInfo:
        """Construct an `ObjectInfo` from a `TypeInfo` instance."""
        parent_type_key: str | None = None
        if type_info.parent_type_info is not None:
            parent_type_key = type_info.parent_type_info.type_key

        # Detect __ffi_init__ from TypeMethod or TypeAttrColumn.
        has_init = any(m.name == "__ffi_init__" for m in type_info.methods)
        if not has_init:
            has_init = _lookup_type_attr(type_info.type_index, "__ffi_init__") is not None

        # Walk parent chain (parent-first) to collect all init-eligible fields.
        init_fields: list[InitFieldInfo] = []
        if has_init:
            ti: TypeInfo | None = type_info
            chain: list[TypeInfo] = []
            while ti is not None:
                chain.append(ti)
                ti = ti.parent_type_info
            for ancestor_info in reversed(chain):
                for field in ancestor_info.fields:
                    if not field.c_init:
                        continue
                    init_fields.append(
                        InitFieldInfo(
                            name=field.name,
                            schema=NamedTypeSchema(
                                name=field.name,
                                schema=_parse_type_schema(field.metadata["type_schema"]),
                            ),
                            kw_only=field.c_kw_only,
                            has_default=field.c_has_default,
                        )
                    )

        return ObjectInfo(
            fields=[
                NamedTypeSchema(
                    name=field.name,
                    schema=_parse_type_schema(field.metadata["type_schema"]),
                )
                for field in type_info.fields
            ],
            methods=[
                FuncInfo(
                    schema=NamedTypeSchema(
                        name=C.FN_NAME_MAP.get(method.name, method.name),
                        schema=_parse_type_schema(method.metadata["type_schema"]),
                    ),
                    is_member=not method.is_static,
                )
                for method in type_info.methods
            ],
            type_key=type_info.type_key,
            parent_type_key=parent_type_key,
            init_fields=init_fields,
            has_init=has_init,
        )

    def gen_fields(self, ty_map: Callable[[str], str], indent: int) -> list[str]:
        """Generate field definitions for this object."""
        indent_str = " " * indent
        return [f"{indent_str}{field.name}: {field.repr(ty_map)}" for field in self.fields]

    def gen_methods(self, ty_map: Callable[[str], str], indent: int) -> list[str]:
        """Generate method definitions for this object."""
        indent_str = " " * indent
        ret = []
        for method in self.methods:
            func_name = method.schema.name.rsplit(".", 1)[-1]
            if func_name == "__ffi_init__":
                # __ffi_init__ is installed as an instance method (self, *args, **kwargs) -> None
                # by _install_ffi_init_attr, regardless of the C++ static registration.
                ret.append(self._gen_ffi_init_from_method(method, ty_map, indent))
                continue
            if not method.is_member:
                ret.append(f"{indent_str}@staticmethod")
            ret.append(method.gen(ty_map, indent))
        return ret

    @staticmethod
    def _gen_ffi_init_from_method(
        method: FuncInfo, ty_map: Callable[[str], str], indent: int
    ) -> str:
        """Render ``__ffi_init__`` TypeMethod as an instance method returning None."""
        indent_str = " " * indent
        schema = method.schema
        # Subclass __ffi_init__ signatures legitimately differ from the parent
        # (different fields → different constructor params), so suppress LSP.
        ignore = "  # ty: ignore[invalid-method-override]"
        if schema.origin != "Callable" or not schema.args:
            ty_map("Any")
            return f"{indent_str}def __ffi_init__(self, *args: Any) -> None: ...{ignore}"
        # schema.args[0] is return type, schema.args[1:] are param types.
        parts: list[str] = []
        for i, arg in enumerate(schema.args[1:]):
            parts.append(f"_{i}: {arg.repr(ty_map)}")
        if parts:
            params = ", ".join(parts)
            return f"{indent_str}def __ffi_init__(self, {params}, /) -> None: ...{ignore}"
        return f"{indent_str}def __ffi_init__(self) -> None: ...{ignore}"

    def gen_ffi_init(self, ty_map: Callable[[str], str], indent: int) -> list[str]:
        """Generate a ``__ffi_init__`` stub when it's not already in TypeMethod.

        For types whose ``__ffi_init__`` is auto-generated by ``RegisterFFIInit``
        (TypeAttrColumn only), synthesize a static-method stub from field metadata.
        Types that already have ``__ffi_init__`` in TypeMethod (from explicit
        ``refl::init<>``) get it via ``gen_methods`` instead.
        """
        if not self.has_init:
            return []
        # If __ffi_init__ is already in methods (from TypeMethod), gen_methods handles it.
        if any(m.schema.name.rsplit(".", 1)[-1] == "__ffi_init__" for m in self.methods):
            return []
        return self._gen_ffi_init_from_fields(ty_map, indent)

    def gen_init(self, ty_map: Callable[[str], str], indent: int) -> list[str]:
        """Generate an ``__init__`` stub from init-eligible field metadata."""
        if not self.has_init:
            return []
        return self._gen_init_from_fields(ty_map, indent)

    def _format_field_params(self, ty_map: Callable[[str], str]) -> str:
        """Format init-eligible fields as a parameter string with defaults and kw_only."""
        positional = [f for f in self.init_fields if not f.kw_only]
        kw_only = [f for f in self.init_fields if f.kw_only]

        pos_required = [f for f in positional if not f.has_default]
        pos_default = [f for f in positional if f.has_default]
        kw_required = [f for f in kw_only if not f.has_default]
        kw_default = [f for f in kw_only if f.has_default]

        parts: list[str] = []
        for f in pos_required:
            parts.append(f"{f.name}: {f.schema.repr(ty_map)}")
        for f in pos_default:
            parts.append(f"{f.name}: {f.schema.repr(ty_map)} = ...")
        if kw_required or kw_default:
            parts.append("*")
            for f in kw_required:
                parts.append(f"{f.name}: {f.schema.repr(ty_map)}")
            for f in kw_default:
                parts.append(f"{f.name}: {f.schema.repr(ty_map)} = ...")

        return ", ".join(parts)

    def _gen_init_from_fields(self, ty_map: Callable[[str], str], indent: int) -> list[str]:
        """Generate ``__init__`` from init-eligible field metadata (auto-generated init)."""
        indent_str = " " * indent
        params = self._format_field_params(ty_map)
        if params:
            return [f"{indent_str}def __init__(self, {params}) -> None: ..."]
        return [f"{indent_str}def __init__(self) -> None: ..."]

    def _gen_ffi_init_from_fields(self, ty_map: Callable[[str], str], indent: int) -> list[str]:
        """Generate ``__ffi_init__`` stub from field metadata for auto-generated init."""
        indent_str = " " * indent
        # Subclass __ffi_init__ signatures legitimately differ from the parent
        # (different fields → different constructor params), so suppress LSP.
        ignore = "  # ty: ignore[invalid-method-override]"
        params = self._format_field_params(ty_map)
        if params:
            return [f"{indent_str}def __ffi_init__(self, {params}) -> None: ...{ignore}"]
        return [f"{indent_str}def __ffi_init__(self) -> None: ...{ignore}"]
