"""Module builder — assemble BuilderSpecs into complete ModuleNode IR.

Collects imports, normalizes type annotations, and auto-discovers
TYPE_CHECKING imports so pyright can resolve every annotation.
"""

from __future__ import annotations

from code_ir import ModuleNode

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


def specs_to_ir_module(specs: list[BuilderSpec], *, manifest: dict | None = None) -> ModuleNode:
    """Convert a list of BuilderSpecs into a ModuleNode IR.

    Collects imports from all specs, normalizes type annotations, and
    auto-discovers TYPE_CHECKING imports so pyright can resolve every annotation.
    """
    all_import_lines: list[str] = [
        "from __future__ import annotations",
    ]
    optional_import_tuples: list[tuple[str, str]] = []
    for spec in specs:
        all_import_lines.extend(gen_runtime_imports(spec))
        opt = gen_optional_import(spec)
        if opt is not None:
            optional_import_tuples.append(opt)

    classes = [spec_to_ir(spec) for spec in specs]

    # Normalize type annotations (same as stubs: types.X → X, unresolvable → Any)
    extra_imports = normalize_stub_classes(classes)
    for imp in extra_imports:
        if imp in OPTIONAL_MODULE_IMPORTS:
            optional_import_tuples.append((imp, "pass"))
        else:
            all_import_lines.append(imp)

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

    # Build set of already-imported names
    already_imported: set[str] = {"BuilderBase"}
    builder_names = {cls.name for cls in classes}
    for spec in specs:
        if not spec.is_composite and not spec.is_standalone:
            already_imported.add(adk_import_name(spec))

    # Collect TYPE_CHECKING-guarded imports for type annotations
    tc_import_lines: list[str] = []

    # 1. Additional stdlib typing imports (Union, Literal, Mapping, TextIO, 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()):
        tc_import_lines.append(f"from {module} import {', '.join(sorted(names))}")

    # 2. 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
            if name in builder_names and f"_ADK_{name}" in already_imported:
                continue
            if name in type_map:
                tc_import_lines.append(type_map[name])
                already_imported.add(name)

    # Add optional imports to TYPE_CHECKING as well, so pyright resolves the
    # real type for return annotations (the runtime try/except fallback sets
    # the name to None which pyright rejects in type expressions).
    for import_line, _fallback in optional_import_tuples:
        tc_import_lines.append(import_line)

    if tc_import_lines:
        all_import_lines.append("from typing import TYPE_CHECKING")

    return ModuleNode(
        doc="Auto-generated by adk-fluent generator. Manual edits will be overwritten.",
        imports=all_import_lines,
        classes=classes,
        type_checking_imports=tc_import_lines,
        optional_imports=optional_import_tuples,
    )
