"""Orchestration for extracting all four context types."""

from __future__ import annotations

import ast
import logging
from typing import TYPE_CHECKING

import libcst as cst

from .dependencies import (
    collect_top_level_defs_with_dependencies,
    mark_defs_for_functions,
    remove_unused_definitions_by_function_names,
)
from .imports import (
    add_needed_imports_from_module,
    gather_source_imports,
)
from .models import (
    AllContextResults,
    CodeContextType,
    CodeString,
    CodeStringsMarkdown,
    FileContextCache,
)
from .pruning import parse_code_and_prune_cst

if TYPE_CHECKING:
    from pathlib import Path

    from codeflash_python._model import FunctionSource

log = logging.getLogger(__name__)


def extract_contexts_for_file(  # noqa: C901, PLR0912, PLR0913, PLR0915
    file_path: Path,
    fto_names: set[str],
    hoh_names: set[str],
    rw_helper_fqns: set[str],
    all_helper_fqns: set[str],
    project_root: Path,
) -> tuple[
    CodeString | None,
    CodeString | None,
    CodeString | None,
    CodeString | None,
    FileContextCache | None,
]:
    """Extract context types from a single source file.

    Returns *(read_writable, read_only, hashing, testgen, cache)*.
    *read_writable* is None when *fto_names* is empty
    (helpers-of-helpers-only file).
    """
    try:
        original_code = file_path.read_text("utf8")
    except Exception:  # noqa: BLE001
        log.debug(
            "Error reading %s",
            file_path,
            exc_info=True,
        )
        return None, None, None, None, None

    try:
        original_module = cst.parse_module(original_code)
    except Exception:  # noqa: BLE001
        log.debug(
            "Failed to parse %s with libcst",
            file_path,
            exc_info=True,
        )
        return None, None, None, None, None

    try:
        relative_path = file_path.resolve().relative_to(
            project_root.resolve(),
        )
    except ValueError:
        relative_path = file_path

    all_names = fto_names | hoh_names
    rw_cs: CodeString | None = None
    ro_cs: CodeString | None = None
    hash_cs: CodeString | None = None
    tg_cs: CodeString | None = None

    if fto_names:
        # FTO file: collect deps once, mark for different sets
        base_defs = collect_top_level_defs_with_dependencies(
            original_module,
        )
        fto_defs = mark_defs_for_functions(
            base_defs,
            fto_names,
        )
        rw_cleaned = remove_unused_definitions_by_function_names(
            original_module,
            fto_names,
            defs_with_usages=fto_defs,
        )
        if hoh_names:
            all_defs = mark_defs_for_functions(
                base_defs,
                all_names,
            )
            all_cleaned = remove_unused_definitions_by_function_names(
                original_module,
                all_names,
                defs_with_usages=all_defs,
            )
        else:
            all_cleaned = rw_cleaned
    else:
        # HoH-only file
        all_cleaned = remove_unused_definitions_by_function_names(
            original_module,
            all_names,
        )
        fto_defs = None
        rw_cleaned = None

    # Pre-compute source imports once for this file
    src_gathered = gather_source_imports(
        original_module,
        file_path,
        project_root,
    )

    # READ_WRITABLE (FTO files only)
    if fto_names and rw_cleaned is not None:
        try:
            rw_pruned = parse_code_and_prune_cst(
                rw_cleaned,
                CodeContextType.READ_WRITABLE,
                fto_names,
                set(),
                remove_docstrings=False,
                defs_with_usages=fto_defs,
            )
            if rw_pruned.code.strip():
                rw_cs = CodeString(
                    code=add_needed_imports_from_module(
                        src_module_code=original_module,
                        dst_module_code=rw_pruned,
                        src_path=file_path,
                        dst_path=file_path,
                        project_root=project_root,
                        helper_fqns=rw_helper_fqns,
                        gathered_imports=src_gathered,
                    ),
                    file_path=relative_path,
                )
        except ValueError:
            log.debug(
                "Error extracting RW context",
                exc_info=True,
            )

    # READ_ONLY
    ro_pruned_code_str: str | None = None
    try:
        ro_pruned = parse_code_and_prune_cst(
            all_cleaned,
            CodeContextType.READ_ONLY,
            fto_names,
            hoh_names,
            remove_docstrings=False,
        )
        ro_pruned_code_str = ro_pruned.code.strip()
        if ro_pruned_code_str:
            ro_cs = CodeString(
                code=add_needed_imports_from_module(
                    src_module_code=original_module,
                    dst_module_code=ro_pruned,
                    src_path=file_path,
                    dst_path=file_path,
                    project_root=project_root,
                    helper_fqns=all_helper_fqns,
                    gathered_imports=src_gathered,
                ),
                file_path=relative_path,
            )
    except ValueError:
        log.debug(
            "Error extracting RO context",
            exc_info=True,
        )

    # HASHING
    try:
        hash_pruned = parse_code_and_prune_cst(
            all_cleaned,
            CodeContextType.HASHING,
            fto_names,
            hoh_names,
            remove_docstrings=True,
        )
        if hash_pruned.code.strip():
            hash_cs = CodeString(
                code=ast.unparse(ast.parse(hash_pruned.code)),
                file_path=relative_path,
            )
    except ValueError:
        log.debug(
            "Error extracting HASHING context",
            exc_info=True,
        )

    # TESTGEN — reuse RO result when pruned code is identical
    try:
        tg_pruned = parse_code_and_prune_cst(
            all_cleaned,
            CodeContextType.TESTGEN,
            fto_names,
            hoh_names,
            remove_docstrings=False,
        )
        tg_pruned_code_str = tg_pruned.code.strip()
        if tg_pruned_code_str:
            if (
                ro_cs is not None
                and ro_pruned_code_str is not None
                and tg_pruned_code_str == ro_pruned_code_str
            ):
                tg_cs = CodeString(
                    code=ro_cs.code,
                    file_path=relative_path,
                )
            else:
                tg_cs = CodeString(
                    code=add_needed_imports_from_module(
                        src_module_code=original_module,
                        dst_module_code=tg_pruned,
                        src_path=file_path,
                        dst_path=file_path,
                        project_root=project_root,
                        helper_fqns=all_helper_fqns,
                        gathered_imports=src_gathered,
                    ),
                    file_path=relative_path,
                )
    except ValueError:
        log.debug(
            "Error extracting TESTGEN context",
            exc_info=True,
        )

    cache = FileContextCache(
        original_module=original_module,
        cleaned_module=all_cleaned,
        fto_names=fto_names,
        hoh_names=hoh_names,
        helper_fqns=all_helper_fqns,
        file_path=file_path,
        relative_path=relative_path,
        gathered_imports=src_gathered,
    )
    return rw_cs, ro_cs, hash_cs, tg_cs, cache


def extract_all_contexts(  # noqa: C901, PLR0912
    helpers_of_fto: dict[Path, set[FunctionSource]],
    helpers_of_helpers: dict[Path, set[FunctionSource]],
    project_root: Path,
) -> AllContextResults:
    """Extract and combine all four context types.

    Processes each file once (single CST parse), then
    prunes into the four context types with appropriate
    imports.
    """
    # Deduplicate: remove HoH entries overlapping with FTO
    hoh_deduped: dict[Path, set[FunctionSource]] = {}
    hoh_no_overlap: dict[Path, set[FunctionSource]] = {}
    for fp, sources in helpers_of_helpers.items():
        if fp in helpers_of_fto:
            diff = sources - helpers_of_fto[fp]
            if diff:
                hoh_deduped[fp] = diff
        else:
            hoh_no_overlap[fp] = sources

    rw_strings: list[CodeString] = []
    ro_strings: list[CodeString] = []
    hash_strings: list[CodeString] = []
    tg_strings: list[CodeString] = []
    all_caches: list[FileContextCache] = []

    # Files containing FTO helpers (all 4 contexts)
    for fp, fto_sources in helpers_of_fto.items():
        fto_names = {f.qualified_name for f in fto_sources}
        hoh_funcs = hoh_deduped.get(fp, set())
        hoh_names = {f.qualified_name for f in hoh_funcs}
        rw_fqns = {f.fully_qualified_name for f in fto_sources}
        all_fqns = rw_fqns | {f.fully_qualified_name for f in hoh_funcs}

        rw, ro, hsh, tg, cache = extract_contexts_for_file(
            fp,
            fto_names,
            hoh_names,
            rw_fqns,
            all_fqns,
            project_root,
        )
        if rw is not None:
            rw_strings.append(rw)
        if ro is not None:
            ro_strings.append(ro)
        if hsh is not None:
            hash_strings.append(hsh)
        if tg is not None:
            tg_strings.append(tg)
        if cache is not None:
            all_caches.append(cache)

    # HoH-only files (RO/HASH/TESTGEN only)
    for fp, hoh_sources in hoh_no_overlap.items():
        hoh_names = {f.qualified_name for f in hoh_sources}
        hoh_fqns = {f.fully_qualified_name for f in hoh_sources}

        _, ro, hsh, tg, cache = extract_contexts_for_file(
            fp,
            set(),
            hoh_names,
            set(),
            hoh_fqns,
            project_root,
        )
        if ro is not None:
            ro_strings.append(ro)
        if hsh is not None:
            hash_strings.append(hsh)
        if tg is not None:
            tg_strings.append(tg)
        if cache is not None:
            all_caches.append(cache)

    return AllContextResults(
        read_writable=CodeStringsMarkdown(code_strings=rw_strings),
        read_only=CodeStringsMarkdown(code_strings=ro_strings),
        hashing=CodeStringsMarkdown(code_strings=hash_strings),
        testgen=CodeStringsMarkdown(code_strings=tg_strings),
        file_caches=all_caches,
    )
