"""Check for missing or extra exports in umbrella package.

Compares __all__ exports from submodules with umbrella package exports
to identify discrepancies. Helps maintain API consistency across packages.
"""

import ast
from pathlib import Path
from typing import NamedTuple, cast


class ExportDiff(NamedTuple):
    """Difference in exports between submodule and umbrella package."""

    submodule: str
    missing_in_umbrella: set[str]  # In submodule but not in umbrella
    extra_in_umbrella: set[str]  # In umbrella but not in submodule


def extract_all_from_init(init_file: Path) -> set[str]:
    """Extract __all__ exports from an __init__.py file.

    Args:
        init_file: Path to __init__.py file

    Returns:
        Set of exported names from __all__
    """
    if not init_file.exists():
        return set()

    try:
        with open(init_file, encoding="utf-8") as f:
            tree = ast.parse(f.read())

        for node in ast.walk(tree):
            if isinstance(node, ast.Assign):
                for target in node.targets:
                    if isinstance(target, ast.Name) and target.id == "__all__":
                        if isinstance(node.value, (ast.List, ast.Tuple)):
                            return {
                                cast(str, elt.value)
                                for elt in node.value.elts
                                if isinstance(elt, ast.Constant)
                            }
    except Exception as e:
        print(f"Warning: Could not parse {init_file}: {e}")

    return set()


def extract_imports_from_init(init_file: Path, submodule_name: str) -> set[str]:
    """Extract what umbrella imports from a specific submodule.

    Args:
        init_file: Path to umbrella __init__.py
        submodule_name: Name of submodule (e.g., 'questra_data')

    Returns:
        Set of names imported from the submodule
    """
    if not init_file.exists():
        return set()

    try:
        with open(init_file, encoding="utf-8") as f:
            tree = ast.parse(f.read())

        imports = set()
        for node in ast.walk(tree):
            if isinstance(node, ast.ImportFrom):
                if node.module == submodule_name:
                    for alias in node.names:
                        imports.add(alias.name)

        return imports
    except Exception as e:
        print(f"Warning: Could not parse {init_file}: {e}")

    return set()


def check_submodule_exports(
    submodule_name: str, submodule_init: Path, umbrella_init: Path
) -> ExportDiff | None:
    """Check exports for a single submodule.

    Args:
        submodule_name: Name of submodule package (e.g., 'questra_data')
        submodule_init: Path to submodule's __init__.py
        umbrella_init: Path to umbrella's __init__.py

    Returns:
        ExportDiff if there are differences, None otherwise
    """
    submodule_exports = extract_all_from_init(submodule_init)
    umbrella_imports = extract_imports_from_init(umbrella_init, submodule_name)

    # What's in submodule but not imported/exported by umbrella
    missing = submodule_exports - umbrella_imports

    # What umbrella imports from submodule but might be outdated
    # (not in submodule's __all__ anymore)
    extra = umbrella_imports - submodule_exports

    if missing or extra:
        return ExportDiff(
            submodule=submodule_name,
            missing_in_umbrella=missing,
            extra_in_umbrella=extra,
        )

    return None


def print_diff(diff: ExportDiff) -> None:
    """Pretty-print an export diff.

    Args:
        diff: Export difference to display
    """
    print(f"\n{'=' * 70}")
    print(f"Submodule: {diff.submodule}")
    print(f"{'=' * 70}")

    if diff.missing_in_umbrella:
        print(f"\n[!] Missing in umbrella (exported by {diff.submodule}):")
        for name in sorted(diff.missing_in_umbrella):
            print(f"   - {name}")

    if diff.extra_in_umbrella:
        print(f"\n[?] In umbrella but not in {diff.submodule}.__all__:")
        for name in sorted(diff.extra_in_umbrella):
            print(f"   - {name}")


def check_all_exports() -> list[ExportDiff]:
    """Check exports for all submodules.

    Returns:
        List of export differences found
    """
    packages_dir = Path(__file__).parent.parent.parent
    umbrella_init = packages_dir / "questra" / "src" / "questra" / "__init__.py"

    submodules = [
        (
            "questra_data",
            packages_dir / "data" / "src" / "questra_data" / "__init__.py",
        ),
        (
            "questra_authentication",
            packages_dir
            / "authentication"
            / "src"
            / "questra_authentication"
            / "__init__.py",
        ),
    ]

    diffs = []
    for submodule_name, submodule_init in submodules:
        if not submodule_init.exists():
            print(f"Warning: Submodule init not found: {submodule_init}")
            continue

        diff = check_submodule_exports(submodule_name, submodule_init, umbrella_init)
        if diff:
            diffs.append(diff)

    return diffs


def main() -> None:
    """Run export checks and display results."""
    print("Checking exports between submodules and umbrella package...")

    diffs = check_all_exports()

    if not diffs:
        print("\n[OK] All exports are in sync!")
        return

    print(f"\n[!] Found {len(diffs)} submodule(s) with export differences:")

    for diff in diffs:
        print_diff(diff)

    print(f"\n{'=' * 70}")
    print("Review these differences and update src/questra/__init__.py manually.")
    print("Consider:")
    print("  - Add missing exports if they should be public API")
    print("  - Remove extra exports if they're deprecated")
    print("  - Keep differences if intentional (internal APIs)")
    print(f"{'=' * 70}\n")


if __name__ == "__main__":
    main()
