"""Stub generation — produce .pyi type stubs from BuilderSpecs.

Reuses spec_to_ir() since emit_stub() already ignores method bodies.
Handles import resolution, type normalization, and name conflict resolution
specifically for the stub output target.
"""

from __future__ import annotations

import hashlib
import json

from code_ir import ModuleNode

from .imports import TYPING_NAMES, adk_import_name, build_type_import_map
from .ir_builders import spec_to_ir
from .spec import BuilderSpec
from .type_normalization import (
    collect_stub_type_refs,
    normalize_stub_classes,
    resolve_stub_name_conflicts,
)


def specs_to_ir_stub_module(specs: list[BuilderSpec], adk_version: str, *, manifest: dict | None = None) -> ModuleNode:
    """Build a ModuleNode suitable for stub (.pyi) emission.

    Reuses spec_to_ir() since emit_stub() already ignores method bodies
    and just outputs ``...`` for each method.

    When *manifest* is provided, the emitted imports are auto-discovered
    from type references in method signatures so that pyright/mypy can
    resolve every annotation.
    """
    classes = [spec_to_ir(spec) for spec in specs]

    # --- Pre-processing: normalize types in method signatures ---
    extra_imports = normalize_stub_classes(classes)

    # Collect all type names referenced in method signatures
    refs = collect_stub_type_refs(classes)

    # --- Build imports based on what's actually referenced ---
    all_import_lines: list[str] = list(extra_imports)
    already_imported: set[str] = set()

    # 1. Base stdlib typing imports — only include names that are referenced
    _BASE_COLLECTIONS = {"AsyncGenerator", "Callable"}
    _BASE_TYPING = {"Any", "Self"}
    needed_collections = sorted(_BASE_COLLECTIONS & refs)
    needed_typing = sorted(_BASE_TYPING & refs)
    if needed_collections:
        all_import_lines.append(f"from collections.abc import {', '.join(needed_collections)}")
        already_imported.update(needed_collections)
    if needed_typing:
        all_import_lines.append(f"from typing import {', '.join(needed_typing)}")
        already_imported.update(needed_typing)

    # 2. Internal import — always needed for base class
    all_import_lines.append("from adk_fluent._base import BuilderBase")
    already_imported.add("BuilderBase")

    # 3. ADK source class imports (build targets)
    builder_names = {cls.name for cls in classes}
    for spec in specs:
        if not spec.is_composite and not spec.is_standalone:
            module_path = ".".join(spec.source_class.split(".")[:-1])
            class_name = spec.source_class.split(".")[-1]
            import_name = adk_import_name(spec)
            if import_name != class_name:
                all_import_lines.append(f"from {module_path} import {class_name} as {import_name}")
            else:
                all_import_lines.append(f"from {module_path} import {class_name}")
            already_imported.add(import_name)

    # 4. Additional stdlib typing imports (Union, Optional, Literal, etc.)
    typing_needed: dict[str, list[str]] = {}
    for name in sorted(refs):
        if name in already_imported:
            continue
        if name in TYPING_NAMES:
            module = TYPING_NAMES[name]
            typing_needed.setdefault(module, []).append(name)
            already_imported.add(name)
    for module, names in sorted(typing_needed.items()):
        all_import_lines.append(f"from {module} import {', '.join(sorted(names))}")

    # 5. Domain type imports (ADK/genai types) from manifest
    if manifest:
        type_map = build_type_import_map(manifest)
        for name in sorted(refs):
            if name in already_imported:
                continue
            # Skip types that conflict with builder class names — these
            # are already imported as _ADK_ aliases via build targets.
            if name in builder_names and f"_ADK_{name}" in already_imported:
                continue
            if name in type_map:
                all_import_lines.append(type_map[name])
                already_imported.add(name)

    # 6. Resolve name conflicts: update method signatures to use _ADK_ prefix
    # where a bare type name would shadow a builder class definition.
    resolve_stub_name_conflicts(classes, already_imported)

    # Content hash of inputs for provenance — deterministic, not time-based
    manifest_hash = ""
    if manifest is not None:
        raw = json.dumps(manifest, sort_keys=True).encode()
        manifest_hash = f" (manifest {hashlib.sha256(raw).hexdigest()[:12]})"

    header = f"AUTO-GENERATED by adk-fluent generator -- do not edit manually\nGenerated from google-adk {adk_version}{manifest_hash}"

    return ModuleNode(
        doc=header,
        imports=all_import_lines,
        classes=classes,
    )
