#!/usr/bin/env python3
"""
Generate type stubs for ucon.dimension from runtime introspection.

This script imports ucon.dimension and generates a .pyi stub file by:
1. Extracting all registered dimensions from _DIMENSION_ATTRS
2. Extracting all module-level constants from __all__
3. Introspecting the Dimension class for methods and properties
4. Generating complete type annotations

Usage:
    python scripts/generate_dimension_stubs.py > ucon/dimension.pyi
    python scripts/generate_dimension_stubs.py --check  # Verify stub is current
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

# Ensure we import from the local package
sys.path.insert(0, str(Path(__file__).parent.parent))

import ucon.dimension as dim_module
from ucon.dimension import Dimension, _DIMENSION_ATTRS, __all__ as DIM_ALL


HEADER = '''\
# Auto-generated type stubs for ucon.dimension
# Generated by: scripts/generate_dimension_stubs.py
# Regenerate with: make stubs
# DO NOT EDIT MANUALLY

from dataclasses import dataclass
from fractions import Fraction
from typing import ClassVar

from ucon.basis import Basis, BasisTransform, Vector

'''


def generate_module_functions() -> list[str]:
    """Generate module-level function stubs."""
    lines = []

    lines.append("# Registry and resolution")
    lines.append("def resolve(vector: Vector) -> Dimension: ...")
    lines.append("def basis() -> tuple[Dimension, ...]: ...")
    lines.append("def all_dimensions() -> tuple[Dimension, ...]: ...")
    lines.append("")

    return lines


def generate_dimension_class() -> list[str]:
    """Generate the Dimension class stub."""
    lines = []

    lines.append("@dataclass(frozen=True)")
    lines.append("class Dimension:")
    lines.append('    """A physical dimension represented by an exponent vector over a basis."""')
    lines.append("")

    # Dataclass fields
    lines.append("    vector: Vector")
    lines.append("    name: str | None = None")
    lines.append("    symbol: str | None = None")
    lines.append("    tag: str | None = None")
    lines.append("")

    # Class-level attributes from _DIMENSION_ATTRS registry
    lines.append("    # Class attributes for IDE completion (from _DIMENSION_ATTRS)")
    for name in sorted(_DIMENSION_ATTRS.keys()):
        lines.append(f"    {name}: ClassVar[Dimension]")
    lines.append("")

    # Class methods
    lines.append("    @classmethod")
    lines.append("    def from_components(")
    lines.append("        cls,")
    lines.append("        basis: Basis | None = None,")
    lines.append("        *,")
    lines.append("        name: str | None = None,")
    lines.append("        symbol: str | None = None,")
    lines.append("        **components: int | float | Fraction,")
    lines.append("    ) -> Dimension: ...")
    lines.append("")

    lines.append("    @classmethod")
    lines.append("    def pseudo(")
    lines.append("        cls,")
    lines.append("        tag: str,")
    lines.append("        *,")
    lines.append("        name: str | None = None,")
    lines.append("        symbol: str | None = None,")
    lines.append("        basis: Basis | None = None,")
    lines.append("    ) -> Dimension: ...")
    lines.append("")

    # Instance methods
    lines.append("    def in_basis(self, transform: BasisTransform) -> Dimension: ...")
    lines.append("")

    # Properties
    lines.append("    @property")
    lines.append("    def basis(self) -> Basis: ...")
    lines.append("    @property")
    lines.append("    def value(self) -> Vector: ...")
    lines.append("    @property")
    lines.append("    def is_pseudo(self) -> bool: ...")
    lines.append("    @property")
    lines.append("    def is_dimensionless(self) -> bool: ...")
    lines.append("")

    # Methods
    lines.append("    def is_base(self) -> bool: ...")
    lines.append("    def base_expansion(self) -> dict[Dimension, Fraction]: ...")
    lines.append("")

    # Operators
    lines.append("    def __mul__(self, other: Dimension) -> Dimension: ...")
    lines.append("    def __truediv__(self, other: Dimension) -> Dimension: ...")
    lines.append("    def __pow__(self, power: int | float | Fraction) -> Dimension: ...")
    lines.append("    def __eq__(self, other: object) -> bool: ...")
    lines.append("    def __hash__(self) -> int: ...")
    lines.append("    def __bool__(self) -> bool: ...")
    lines.append("    def __repr__(self) -> str: ...")
    lines.append("")

    return lines


def generate_module_constants() -> list[str]:
    """Generate module-level constant stubs from __all__."""
    lines = []

    lines.append("# Module constants")

    # Extract dimension constants (uppercase names)
    for name in DIM_ALL:
        obj = getattr(dim_module, name, None)
        if obj is None:
            continue

        if isinstance(obj, Dimension):
            lines.append(f"{name}: Dimension")
        elif name == "SI":
            lines.append(f"{name}: Basis")
        # Skip functions and classes (already handled)

    lines.append("")
    return lines


def generate() -> str:
    """Generate the complete stub file."""
    lines = [HEADER]

    lines.extend(generate_module_functions())
    lines.extend(generate_dimension_class())
    lines.extend(generate_module_constants())

    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description="Generate dimension.pyi stubs")
    parser.add_argument(
        "--check",
        action="store_true",
        help="Check if existing stub matches generated (exit 1 if different)",
    )
    parser.add_argument(
        "--output", "-o",
        type=Path,
        default=None,
        help="Output file (default: stdout)",
    )
    args = parser.parse_args()

    generated = generate()

    if args.check:
        stub_path = Path(__file__).parent.parent / "ucon" / "dimension.pyi"
        if not stub_path.exists():
            print(f"ERROR: {stub_path} does not exist", file=sys.stderr)
            sys.exit(1)

        existing = stub_path.read_text()
        if existing.strip() != generated.strip():
            print("ERROR: dimension.pyi is out of date. Run 'make stubs' to regenerate.", file=sys.stderr)
            sys.exit(1)

        print("OK: dimension.pyi is up to date")
        sys.exit(0)

    if args.output:
        args.output.write_text(generated)
        print(f"Wrote {args.output}", file=sys.stderr)
    else:
        print(generated)


if __name__ == "__main__":
    main()
