"""Testgen context enrichment via class resolution and init stub extraction.

Resolves imported classes in testgen context to their definitions (project
classes get full source, third-party classes get __init__ stubs). Handles
dataclasses, attrs, and NamedTuple synthetic constructors.
"""

from __future__ import annotations

import ast
import logging
import os
from typing import TYPE_CHECKING

from ._ast_helpers import (
    ImportCollector,
    collect_existing_class_names,
    collect_import_aliases,
    collect_type_names_from_annotation,
    find_class_node_by_name,
    get_expr_name,
)
from ._class_analysis import (
    build_import_from_map,
    build_synthetic_init_stub,
    collect_synthetic_constructor_type_names,
    get_class_start_line,
    resolve_instance_class_name,
    should_use_raw_project_class_context,
)
from .helpers import is_project_path
from .models import CodeString, CodeStringsMarkdown
from .resolve import get_jedi_project

if TYPE_CHECKING:
    from pathlib import Path

    from .._model import FunctionToOptimize


log = logging.getLogger(__name__)

BUILTIN_AND_TYPING_NAMES = frozenset(
    {
        "int",
        "str",
        "float",
        "bool",
        "bytes",
        "bytearray",
        "complex",
        "list",
        "dict",
        "set",
        "frozenset",
        "tuple",
        "type",
        "object",
        "None",
        "NoneType",
        "Ellipsis",
        "NotImplemented",
        "memoryview",
        "range",
        "slice",
        "property",
        "classmethod",
        "staticmethod",
        "super",
        "Optional",
        "Union",
        "Any",
        "List",
        "Dict",
        "Set",
        "FrozenSet",
        "Tuple",
        "Type",
        "Callable",
        "Iterator",
        "Generator",
        "Coroutine",
        "AsyncGenerator",
        "AsyncIterator",
        "Iterable",
        "AsyncIterable",
        "Sequence",
        "MutableSequence",
        "Mapping",
        "MutableMapping",
        "Collection",
        "Awaitable",
        "Literal",
        "Final",
        "ClassVar",
        "TypeVar",
        "TypeAlias",
        "ParamSpec",
        "Concatenate",
        "Annotated",
        "TypeGuard",
        "Self",
        "Unpack",
        "TypeVarTuple",
        "Never",
        "NoReturn",
        "SupportsInt",
        "SupportsFloat",
        "SupportsComplex",
        "SupportsBytes",
        "SupportsAbs",
        "SupportsRound",
        "IO",
        "TextIO",
        "BinaryIO",
        "Pattern",
        "Match",
    }
)


def extract_function_stub_snippet(
    fn_node: ast.FunctionDef | ast.AsyncFunctionDef,
    module_lines: list[str],
) -> str:
    """Extract source lines of a function (including decorators)."""
    start_line = (
        min(d.lineno for d in fn_node.decorator_list)
        if fn_node.decorator_list
        else fn_node.lineno
    )
    return "\n".join(module_lines[start_line - 1 : fn_node.end_lineno])


def extract_init_stub_from_class(  # noqa: C901
    class_name: str,
    module_source: str,
    module_tree: ast.Module,
) -> str | None:
    """Extract or synthesize an ``__init__`` stub for *class_name*.

    For declarative classes (dataclass, attrs, NamedTuple) that lack an
    explicit ``__init__``, a synthetic stub is generated. Otherwise the
    existing ``__init__``, ``__post_init__``, and property methods are
    extracted.
    """
    class_node = find_class_node_by_name(class_name, module_tree)
    if class_node is None:
        return None

    lines = module_source.splitlines()
    import_aliases = collect_import_aliases(module_tree)
    explicit_init_nodes: list[ast.FunctionDef | ast.AsyncFunctionDef] = []
    support_nodes: list[ast.FunctionDef | ast.AsyncFunctionDef] = []
    for item in class_node.body:
        if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
            if item.name == "__init__":
                explicit_init_nodes.append(item)
                support_nodes.append(item)
                continue
            if item.name == "__post_init__":
                support_nodes.append(item)
                continue
            for d in item.decorator_list:
                if (isinstance(d, ast.Name) and d.id == "property") or (
                    isinstance(d, ast.Attribute) and d.attr == "property"
                ):
                    support_nodes.append(item)
                    break

    snippets: list[str] = []
    if not explicit_init_nodes:
        synthetic_init = build_synthetic_init_stub(
            class_node, module_source, import_aliases
        )
        if synthetic_init is not None:
            snippets.append(synthetic_init)
    snippets.extend(
        extract_function_stub_snippet(fn_node, lines)
        for fn_node in support_nodes
    )

    if not snippets:
        return None

    return f"class {class_name}:\n" + "\n".join(snippets)


def extract_imports_for_class(  # noqa: C901, PLR0912
    module_tree: ast.Module,
    class_node: ast.ClassDef,
    module_source: str,
) -> str:
    """Extract import statements needed by *class_node*."""
    needed_names: set[str] = set()

    for base in class_node.bases:
        if isinstance(base, ast.Name):
            needed_names.add(base.id)
        elif isinstance(base, ast.Attribute) and isinstance(
            base.value, ast.Name
        ):
            needed_names.add(base.value.id)

    for decorator in class_node.decorator_list:
        if isinstance(decorator, ast.Name):
            needed_names.add(decorator.id)
        elif isinstance(decorator, ast.Call):
            if isinstance(decorator.func, ast.Name):
                needed_names.add(decorator.func.id)
            elif isinstance(decorator.func, ast.Attribute) and isinstance(
                decorator.func.value, ast.Name
            ):
                needed_names.add(decorator.func.value.id)

    for item in class_node.body:
        if isinstance(item, ast.AnnAssign) and item.annotation:
            needed_names |= collect_type_names_from_annotation(item.annotation)
        elif (
            isinstance(item, ast.Assign)
            and isinstance(item.value, ast.Call)
            and isinstance(item.value.func, ast.Name)
        ):
            needed_names.add(item.value.func.id)

    import_lines: list[str] = []
    source_lines = module_source.split("\n")
    added_imports: set[int] = set()
    for node in module_tree.body:
        if (
            not isinstance(node, (ast.Import, ast.ImportFrom))
            or node.lineno in added_imports
        ):
            continue
        for alias in node.names:
            name = alias.asname or (
                alias.name.split(".")[0]
                if isinstance(node, ast.Import)
                else alias.name
            )
            if name in needed_names:
                import_lines.append(source_lines[node.lineno - 1])
                added_imports.add(node.lineno)
                break

    return "\n".join(import_lines)


def extract_raw_class_context(
    class_node: ast.ClassDef,
    module_source: str,
    module_tree: ast.Module,
) -> str:
    """Extract full class source with needed imports prepended."""
    class_source = "\n".join(
        module_source.splitlines()[
            get_class_start_line(class_node) - 1 : class_node.end_lineno
        ]
    )
    needed_imports = extract_imports_for_class(
        module_tree, class_node, module_source
    )
    if needed_imports:
        return f"{needed_imports}\n\n{class_source}"
    return class_source


def get_module_source_and_tree(
    module_path: Path,
    module_cache: dict[Path, tuple[str, ast.Module]],
) -> tuple[str, ast.Module] | None:
    """Read and parse *module_path*, using *module_cache*."""
    if module_path in module_cache:
        return module_cache[module_path]
    try:
        module_source = module_path.read_text(encoding="utf-8")
        module_tree = ast.parse(module_source)
    except Exception:  # noqa: BLE001
        return None
    module_cache[module_path] = (module_source, module_tree)
    return module_source, module_tree


def parse_and_collect_imports(
    code_context: CodeStringsMarkdown,
) -> tuple[ast.Module, dict[str, str]] | None:
    """Parse combined code and collect import mappings."""
    all_code = "\n".join(cs.code for cs in code_context.code_strings)
    try:
        tree = ast.parse(all_code)
    except SyntaxError:
        return None
    collector = ImportCollector()
    collector.visit(tree)
    return tree, collector.imported_names


def resolve_imported_class_reference(  # noqa: PLR0911
    base_expr_name: str,
    current_module_tree: ast.Module,
    current_module_path: Path,
    project_root_path: Path,
    module_cache: dict[Path, tuple[str, ast.Module]],
) -> tuple[str, Path] | None:
    """Resolve a base class name to ``(class_name, module_path)``."""
    import jedi  # type: ignore[import-untyped]  # noqa: PLC0415

    import_aliases = collect_import_aliases(current_module_tree)
    class_name = base_expr_name.rsplit(".", 1)[-1]
    if (
        "." not in base_expr_name
        and find_class_node_by_name(class_name, current_module_tree)
        is not None
    ):
        return class_name, current_module_path

    resolved_name = base_expr_name
    if base_expr_name in import_aliases:
        resolved_name = import_aliases[base_expr_name]
    elif "." in base_expr_name:
        head, tail = base_expr_name.split(".", 1)
        if head in import_aliases:
            resolved_name = f"{import_aliases[head]}.{tail}"

    if "." not in resolved_name:
        return None

    module_name, class_name = resolved_name.rsplit(".", 1)
    try:
        script_code = f"from {module_name} import {class_name}"
        script = jedi.Script(
            script_code,
            project=get_jedi_project(str(project_root_path)),
        )
        definitions = script.goto(
            1,
            len(f"from {module_name} import ") + len(class_name),
            follow_imports=True,
        )
    except Exception:  # noqa: BLE001
        return None

    if not definitions or definitions[0].module_path is None:
        return None
    module_path = definitions[0].module_path
    if not is_project_path(module_path, project_root_path):
        return None
    if get_module_source_and_tree(module_path, module_cache) is None:
        return None
    return class_name, module_path


def append_project_class_context(  # noqa: PLR0913
    class_name: str,
    module_path: Path,
    project_root_path: Path,
    module_cache: dict[Path, tuple[str, ast.Module]],
    existing_class_names: set[str],
    emitted_classes: set[tuple[Path, str]],
    emitted_class_names: set[str],
    code_strings: list[CodeString],
) -> bool:
    """Append full class context for a project class, recursing into bases."""
    module_result = get_module_source_and_tree(module_path, module_cache)
    if module_result is None:
        return False
    module_source, module_tree = module_result
    class_node = find_class_node_by_name(class_name, module_tree)
    if class_node is None:
        return False

    class_key = (module_path, class_name)
    if class_key in emitted_classes or class_name in existing_class_names:
        return True

    for base in class_node.bases:
        base_expr_name = get_expr_name(base)
        if base_expr_name is None:
            continue
        resolved = resolve_imported_class_reference(
            base_expr_name,
            module_tree,
            module_path,
            project_root_path,
            module_cache,
        )
        if resolved is None:
            continue
        base_name, base_module_path = resolved
        if base_name in existing_class_names:
            continue
        append_project_class_context(
            base_name,
            base_module_path,
            project_root_path,
            module_cache,
            existing_class_names,
            emitted_classes,
            emitted_class_names,
            code_strings,
        )

    code_strings.append(
        CodeString(
            code=extract_raw_class_context(
                class_node, module_source, module_tree
            ),
            file_path=module_path,
        )
    )
    emitted_classes.add(class_key)
    emitted_class_names.add(class_name)
    return True


def collect_type_names_from_function(  # noqa: C901, PLR0912
    func_node: ast.FunctionDef | ast.AsyncFunctionDef,
    tree: ast.Module,
    class_name: str | None,
) -> set[str]:
    """Collect type names from function annotations and isinstance calls."""
    type_names: set[str] = set()
    for arg in (
        func_node.args.args
        + func_node.args.posonlyargs
        + func_node.args.kwonlyargs
    ):
        type_names |= collect_type_names_from_annotation(arg.annotation)
    if func_node.args.vararg:
        type_names |= collect_type_names_from_annotation(
            func_node.args.vararg.annotation
        )
    if func_node.args.kwarg:
        type_names |= collect_type_names_from_annotation(
            func_node.args.kwarg.annotation
        )
    for body_node in ast.walk(func_node):
        if (
            isinstance(body_node, ast.Call)
            and isinstance(body_node.func, ast.Name)
            and body_node.func.id == "isinstance"
            and len(body_node.args) >= 2  # noqa: PLR2004
        ):
            second_arg = body_node.args[1]
            if isinstance(second_arg, ast.Name):
                type_names.add(second_arg.id)
            elif isinstance(second_arg, ast.Tuple):
                for elt in second_arg.elts:
                    if isinstance(elt, ast.Name):
                        type_names.add(elt.id)
        elif isinstance(body_node, ast.Compare) and (
            isinstance(body_node.left, ast.Call)
            and isinstance(body_node.left.func, ast.Name)
            and body_node.left.func.id == "type"
        ):
            for comparator in body_node.comparators:
                if isinstance(comparator, ast.Name):
                    type_names.add(comparator.id)
    if class_name is not None:
        for top_node in ast.walk(tree):
            if (
                isinstance(top_node, ast.ClassDef)
                and top_node.name == class_name
            ):
                for base in top_node.bases:
                    if isinstance(base, ast.Name):
                        type_names.add(base.id)
                break
    return type_names


def build_testgen_context(
    testgen_base: CodeStringsMarkdown,
    project_root_path: Path,
    *,
    include_enrichment: bool = True,
    function_to_optimize: FunctionToOptimize | None = None,
) -> CodeStringsMarkdown:
    """Build enriched testgen context from base extraction.

    Enriches *testgen_base* with resolved class definitions and
    constructor stubs for parameter types used by the function.
    """
    testgen_context = testgen_base

    if include_enrichment:
        enrichment = enrich_testgen_context(
            testgen_context,
            project_root_path,
        )
        if enrichment.code_strings:
            testgen_context = CodeStringsMarkdown(
                code_strings=(
                    testgen_context.code_strings + enrichment.code_strings
                ),
            )

        if function_to_optimize is not None:
            result = parse_and_collect_imports(testgen_context)
            existing_classes = (
                collect_existing_class_names(result[0]) if result else set()
            )
            constructor_stubs = extract_parameter_type_constructors(
                function_to_optimize,
                project_root_path,
                existing_classes,
            )
            if constructor_stubs.code_strings:
                testgen_context = CodeStringsMarkdown(
                    code_strings=(
                        testgen_context.code_strings
                        + constructor_stubs.code_strings
                    ),
                )

    return testgen_context


def enrich_testgen_context(  # noqa: C901, PLR0912, PLR0915
    code_context: CodeStringsMarkdown,
    project_root_path: Path,
) -> CodeStringsMarkdown:
    """Enrich testgen context with resolved class definitions.

    For imported classes in the testgen context, resolve via Jedi to
    their definitions. Project classes get full source, third-party
    classes get ``__init__`` stubs.
    """
    import jedi  # noqa: PLC0415

    result = parse_and_collect_imports(code_context)
    if result is None:
        return CodeStringsMarkdown(code_strings=[])
    tree, imported_names = result

    if not imported_names:
        return CodeStringsMarkdown(code_strings=[])

    existing_classes = collect_existing_class_names(tree)

    code_strings: list[CodeString] = []
    emitted_class_names: set[str] = set()

    extracted_classes: set[tuple[Path, str]] = set()
    module_cache: dict[Path, tuple[str, ast.Module]] = {}

    def extract_class_and_bases(
        cn: str,
        module_path: Path,
        module_source: str,
        module_tree: ast.Module,
    ) -> None:
        """Recursively extract a class and its base classes."""
        if (module_path, cn) in extracted_classes:
            return

        class_node = find_class_node_by_name(cn, module_tree)
        if class_node is None:
            return

        for base in class_node.bases:
            base_name = None
            if isinstance(base, ast.Name):
                base_name = base.id
            elif isinstance(base, ast.Attribute):
                continue

            if base_name and base_name not in existing_classes:
                extract_class_and_bases(
                    base_name,
                    module_path,
                    module_source,
                    module_tree,
                )

        if (module_path, cn) in extracted_classes:
            return

        lines = module_source.split("\n")
        class_source = "\n".join(
            lines[get_class_start_line(class_node) - 1 : class_node.end_lineno]
        )

        code_strings.append(
            CodeString(code=class_source, file_path=module_path)
        )
        extracted_classes.add((module_path, cn))
        emitted_class_names.add(cn)

    for name, module_name in imported_names.items():
        if name in existing_classes or module_name == "__future__":
            continue
        try:
            test_code = f"import {module_name}"
            script = jedi.Script(
                test_code,
                project=get_jedi_project(str(project_root_path)),
            )
            completions = script.goto(1, len(test_code))

            if not completions:
                continue

            module_path = completions[0].module_path
            if not module_path:
                continue

            resolved_module = module_path.resolve()
            module_str = str(resolved_module)
            is_proj = module_str.startswith(
                str(project_root_path.resolve()) + os.sep
            )
            is_third_party = "site-packages" in module_str
            if not is_proj and not is_third_party:
                continue

            mod_result = get_module_source_and_tree(module_path, module_cache)
            if mod_result is None:
                continue
            module_source, module_tree = mod_result

            if is_proj:
                extract_class_and_bases(
                    name,
                    module_path,
                    module_source,
                    module_tree,
                )
                if (
                    module_path,
                    name,
                ) not in extracted_classes:
                    resolved_class = resolve_instance_class_name(
                        name, module_tree
                    )
                    if (
                        resolved_class
                        and resolved_class not in existing_classes
                    ):
                        extract_class_and_bases(
                            resolved_class,
                            module_path,
                            module_source,
                            module_tree,
                        )
            elif is_third_party:
                target_name = name
                if find_class_node_by_name(name, module_tree) is None:
                    resolved_class = resolve_instance_class_name(
                        name, module_tree
                    )
                    if resolved_class:
                        target_name = resolved_class
                if target_name not in emitted_class_names:
                    stub = extract_init_stub_from_class(
                        target_name, module_source, module_tree
                    )
                    if stub:
                        code_strings.append(
                            CodeString(
                                code=stub,
                                file_path=module_path,
                            )
                        )
                        emitted_class_names.add(target_name)

        except Exception:  # noqa: BLE001
            log.debug(
                "Error extracting class definition for %s from %s",
                name,
                module_name,
            )
            continue

    return CodeStringsMarkdown(code_strings=code_strings)


def extract_parameter_type_constructors(  # noqa: C901, PLR0912, PLR0915
    function_to_optimize: FunctionToOptimize,
    project_root_path: Path,
    existing_class_names: set[str],
) -> CodeStringsMarkdown:
    """Extract ``__init__`` stubs for types used in a function's signature.

    Finds types used in annotations and ``isinstance`` checks, then
    resolves them via Jedi. Handles dataclasses, attrs, NamedTuple
    with synthetic ``__init__`` generation. Includes one level of
    transitive extraction for constructor parameter types.
    """
    import jedi  # noqa: PLC0415

    try:
        source = function_to_optimize.file_path.read_text(encoding="utf-8")
        tree = ast.parse(source)
    except Exception:  # noqa: BLE001
        return CodeStringsMarkdown(code_strings=[])

    func_node = None
    for node in ast.walk(tree):
        if (
            isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
            and node.name == function_to_optimize.function_name
        ):
            if (
                function_to_optimize.starting_line is not None
                and node.lineno != function_to_optimize.starting_line
            ):
                continue
            func_node = node
            break
    if func_node is None:
        return CodeStringsMarkdown(code_strings=[])

    type_names = collect_type_names_from_function(
        func_node, tree, function_to_optimize.class_name
    )
    type_names -= BUILTIN_AND_TYPING_NAMES
    type_names -= existing_class_names
    if not type_names:
        return CodeStringsMarkdown(code_strings=[])

    import_map = build_import_from_map(tree)

    code_strings: list[CodeString] = []
    module_cache: dict[Path, tuple[str, ast.Module]] = {}
    emitted_classes: set[tuple[Path, str]] = set()
    emitted_class_names: set[str] = set()

    def append_type_context(
        type_name: str,
        module_name: str,
        *,
        transitive: bool = False,
    ) -> None:
        """Resolve a type and append its class source to context."""
        try:
            script_code = f"from {module_name} import {type_name}"
            script = jedi.Script(
                script_code,
                project=get_jedi_project(str(project_root_path)),
            )
            definitions = script.goto(
                1,
                len(f"from {module_name} import ") + len(type_name),
                follow_imports=True,
            )
            if not definitions:
                return

            module_path = definitions[0].module_path
            if not module_path:
                return
            resolved_module = module_path.resolve()
            module_str = str(resolved_module)
            is_proj = is_project_path(module_path, project_root_path)
            is_tp = "site-packages" in module_str
            if transitive and not is_proj and not is_tp:
                return

            module_result = get_module_source_and_tree(
                module_path, module_cache
            )
            if module_result is None:
                return
            mod_source, mod_tree = module_result

            class_key = (module_path, type_name)
            if (
                class_key in emitted_classes
                or type_name in existing_class_names
            ):
                return

            class_node = find_class_node_by_name(type_name, mod_tree)
            if (
                class_node is not None
                and is_proj
                and should_use_raw_project_class_context(
                    class_node,
                    collect_import_aliases(mod_tree),
                )
                and append_project_class_context(
                    type_name,
                    module_path,
                    project_root_path,
                    module_cache,
                    existing_class_names,
                    emitted_classes,
                    emitted_class_names,
                    code_strings,
                )
            ):
                return

            stub = extract_init_stub_from_class(
                type_name, mod_source, mod_tree
            )
            if stub:
                code_strings.append(
                    CodeString(code=stub, file_path=module_path)
                )
                emitted_classes.add(class_key)
                emitted_class_names.add(type_name)
        except Exception:  # noqa: BLE001
            if transitive:
                log.debug(
                    "Error extracting transitive constructor"
                    " stub for %s from %s",
                    type_name,
                    module_name,
                )
            else:
                log.debug(
                    "Error extracting constructor stub for %s from %s",
                    type_name,
                    module_name,
                )

    for type_name in sorted(type_names):
        module_name = import_map.get(type_name)
        if not module_name:
            continue
        append_type_context(type_name, module_name)

    # Transitive extraction (one level)
    transitive_import_map = dict(import_map)
    for _, cached_tree in module_cache.values():
        for name, module in build_import_from_map(cached_tree).items():
            transitive_import_map.setdefault(name, module)

    emitted_names = (
        type_names
        | existing_class_names
        | emitted_class_names
        | BUILTIN_AND_TYPING_NAMES
    )
    transitive_type_names: set[str] = set()
    for cs in code_strings:
        try:
            stub_tree = ast.parse(cs.code)
        except SyntaxError:
            continue
        ia = collect_import_aliases(stub_tree)
        for stub_node in ast.walk(stub_tree):
            if isinstance(
                stub_node,
                (ast.FunctionDef, ast.AsyncFunctionDef),
            ) and stub_node.name in (
                "__init__",
                "__post_init__",
            ):
                for arg in (
                    stub_node.args.args
                    + stub_node.args.posonlyargs
                    + stub_node.args.kwonlyargs
                ):
                    transitive_type_names |= (
                        collect_type_names_from_annotation(arg.annotation)
                    )
            elif isinstance(stub_node, ast.ClassDef):
                transitive_type_names |= (
                    collect_synthetic_constructor_type_names(stub_node, ia)
                )
    transitive_type_names -= emitted_names
    for type_name in sorted(transitive_type_names):
        module_name = transitive_import_map.get(type_name)
        if not module_name:
            continue
        append_type_context(type_name, module_name, transitive=True)

    return CodeStringsMarkdown(code_strings=code_strings)
