# scripts/generate_stubs.py
import inspect
import os
import sys
from typing import Any, Callable, Type

from QuPRS.pathsum import core as core_module
from QuPRS.pathsum.gates import get_all_gates
from QuPRS.pathsum.gates.base import Gate
from QuPRS.pathsum.gates.builder import build_docstring, build_signature

# Set up the project root and source path for module imports.
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
src_path = os.path.join(project_root, "src")
if src_path not in sys.path:
    sys.path.insert(0, src_path)


def _format_pyi_docstring(doc: str) -> str:
    """
    Format a docstring for inclusion in a .pyi stub file.
    Handles single-line and multi-line docstrings.
    """
    if not doc:
        return '""'
    if "\n" not in doc:
        return f'"""{doc}"""'
    lines = doc.split("\n")
    indented_body = "\n".join([f"    {line}" for line in lines])
    return f'"""\n{indented_body}\n    """'


def _cleanup_type_str(type_str: str) -> str:
    """
    Clean up type annotation strings for stub output.
    Removes module prefixes and standardizes type names.
    """
    if type_str is None:
        return "Any"
    s = (
        str(type_str)
        .replace("QuPRS.pathsum.core.", "")
        .replace("symengine.lib.symengine_wrapper.", "se.")
    )
    return (
        s.replace("sympy.core.expr.", "sp.")
        .replace("typing.", "")
        .replace("<class '", "")
        .replace("'>", "")
    )


def _build_stub_for_member(name: str, member: Any) -> str:
    """
    Generate a .pyi stub string for a class member.
    Handles properties, static/class methods, and regular methods.
    """
    if name.startswith("__") and name not in (
        "__init__",
        "__eq__",
        "__hash__",
        "__repr__",
        "__len__",
        "__getitem__",
    ):
        return ""
    func: Callable[..., Any] | None = None
    decorator = ""
    if isinstance(member, property):
        doc = inspect.getdoc(member.fget) or ""
        return_type = "Any"
        if member.fget and "return" in member.fget.__annotations__:
            return_type = _cleanup_type_str(member.fget.__annotations__["return"])
        return f"\n    @property\n    def {name}(self) -> {return_type}:\n        {_format_pyi_docstring(doc)}\n        ...\n"
    elif isinstance(member, (staticmethod, classmethod)):
        decorator = f"@{type(member).__name__}\n    "
        func = member.__func__
    elif inspect.isfunction(member) or inspect.ismethod(member):
        func = member
    if func:
        try:
            sig = inspect.signature(func)
            doc = inspect.getdoc(func) or ""
            cleaned_sig_str = _cleanup_type_str(str(sig))
            return f"\n    {decorator}def {name}{cleaned_sig_str}:\n        {_format_pyi_docstring(doc)}\n        ...\n"
        except (ValueError, TypeError):
            return f"\n    def {name}(self, *args: Any, **kwargs: Any) -> Any: ...\n"
    return ""


def _build_dynamic_method_stub(gate_cls: Type[Gate]) -> str:
    """
    Generate a .pyi stub for a dynamically injected gate method.
    Used for PathSum gate methods.
    """
    gate_name = gate_cls.gate_name
    signature = build_signature(gate_cls, include_self=True)
    docstring = build_docstring(gate_cls, signature)
    cleaned_sig_str = _cleanup_type_str(str(signature))
    return f"\n    def {gate_name}{cleaned_sig_str} -> 'PathSum':\n        {_format_pyi_docstring(docstring)}\n        ...\n"


def generate_stub_file():
    """
    Main function to generate the .pyi stub file for core classes.
    Writes the output to QuPRS/pathsum/core.pyi.
    """
    print("Starting to generate .pyi stub file...")
    output_path = os.path.join(src_path, "QuPRS", "pathsum", "core.pyi")
    pyi_content = [
        "# This file is auto-generated by scripts/generate_stubs.py.",
        "# Do not edit this file directly.",
        "\nfrom typing import Any, Dict, List, Optional, Set, Tuple, Union",
        "import QuPRS",
        "import symengine as se",
        "import sympy as sp",
        "\n",
    ]
    classes_to_stub = [core_module.Register, core_module.F, core_module.PathSum]
    gate_class_map = get_all_gates()
    gate_names = {gate.gate_name for gate in gate_class_map.values()}
    for cls in classes_to_stub:
        print(f"Processing class: {cls.__name__}...")
        pyi_content.append(f"class {cls.__name__}:")
        class_doc = inspect.getdoc(cls) or ""
        if class_doc:
            pyi_content.append(f"    {_format_pyi_docstring(class_doc)}")
        for name, member in cls.__dict__.items():
            if cls is core_module.PathSum and name in gate_names:
                continue
            pyi_content.append(_build_stub_for_member(name, member))
        if cls is core_module.PathSum:
            pyi_content.append("\n    # --- Dynamically Injected Gate Methods ---")
            for gate_cls in sorted(gate_class_map.values(), key=lambda c: c.gate_name):
                pyi_content.append(_build_dynamic_method_stub(gate_cls))
        pyi_content.append("\n")
    try:
        with open(output_path, "w", encoding="utf-8") as f:
            f.write("\n".join(pyi_content))
        print(f"\nSuccessfully generated stub file at:\n{output_path}")
    except Exception as e:
        print(f"\nError writing to file: {e}")


if __name__ == "__main__":
    generate_stub_file()
