"""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


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 {
                                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
    """
    base_dir = Path(__file__).parent.parent
    umbrella_init = base_dir / "src" / "questra" / "__init__.py"

    submodules = [
        (
            "questra_data",
            base_dir / "data" / "src" / "questra_data" / "__init__.py",
        ),
        (
            "questra_authentication",
            base_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()
