"""Top-level pipeline for context extraction and assembly."""

from __future__ import annotations

import hashlib
import logging
from itertools import chain
from typing import TYPE_CHECKING

from ..analysis._code_utils import find_preexisting_objects
from .enrichment import build_testgen_context
from .fallback import encoded_tokens_len, re_extract_from_cache
from .helpers import discover_helpers_by_names
from .models import (
    CodeContextType,
    CodeOptimizationContext,
    CodeStringsMarkdown,
)
from .orchestration import extract_all_contexts
from .resolve import get_function_source

if TYPE_CHECKING:
    from pathlib import Path

    from .._model import FunctionParent, FunctionToOptimize

log = logging.getLogger(__name__)

OPTIMIZATION_CONTEXT_TOKEN_LIMIT = 64000
TESTGEN_CONTEXT_TOKEN_LIMIT = 64000
READ_WRITABLE_LIMIT_ERROR = (
    "Read-writable code has exceeded token limit, cannot proceed"
)
TESTGEN_LIMIT_ERROR = (
    "Testgen code context has exceeded token limit, cannot proceed"
)


def get_code_optimization_context(  # noqa: C901, PLR0915
    function_to_optimize: FunctionToOptimize,
    project_root: Path,
    optim_token_limit: int = OPTIMIZATION_CONTEXT_TOKEN_LIMIT,
    testgen_token_limit: int = TESTGEN_CONTEXT_TOKEN_LIMIT,
) -> CodeOptimizationContext:
    """Build the full optimization context for a function.

    Discovers helpers two levels deep via Jedi, extracts all four
    context types, applies progressive token-limit fallback, enriches
    testgen context with class definitions, and returns a
    *CodeOptimizationContext* ready for the optimization API.

    Raises *ValueError* if read-writable or testgen context exceeds
    its token limit after all fallback strategies are exhausted.
    """
    # Level 1: discover direct helpers of the function to optimize
    fto_input = {
        function_to_optimize.file_path: {
            function_to_optimize.qualified_name,
        },
    }
    refs_cache: dict[Path, dict[str, list[object]]] = {}
    helpers_of_fto = discover_helpers_by_names(
        fto_input,
        project_root,
        refs_cache=refs_cache,
    )

    # Add the FTO itself into the helpers dict
    fto_source = get_function_source(function_to_optimize, project_root)
    helpers_of_fto.setdefault(
        function_to_optimize.file_path,
        set(),
    ).add(fto_source)

    # Build level 2 input from all helpers of FTO (including FTO)
    level2_input: dict[Path, set[str]] = {}
    for fp, sources in helpers_of_fto.items():
        names: set[str] = set()
        for fs in sources:
            names.add(fs.qualified_name)
            # For class methods, also discover __init__
            if "." in fs.qualified_name:
                class_name = fs.qualified_name.rsplit(".", 1)[0]
                names.add(f"{class_name}.__init__")
        level2_input[fp] = names

    # Level 2: discover transitive helpers (reuse refs cache)
    helpers_of_helpers = discover_helpers_by_names(
        level2_input,
        project_root,
        refs_cache=refs_cache,
    )

    # Collect flat lists for downstream consumers
    fto_source_list = [
        fs for sources in helpers_of_fto.values() for fs in sources
    ]
    hoh_source_list = [
        fs for sources in helpers_of_helpers.values() for fs in sources
    ]

    # Extract all four context types in a single pass per file
    all_ctx = extract_all_contexts(
        helpers_of_fto,
        helpers_of_helpers,
        project_root,
    )

    # Ensure target file is first in RW code blocks
    rw_code = all_ctx.read_writable
    try:
        target_relative = function_to_optimize.file_path.resolve().relative_to(
            project_root.resolve(),
        )
        target_blocks = [
            cs
            for cs in rw_code.code_strings
            if cs.file_path == target_relative
        ]
        other_blocks = [
            cs
            for cs in rw_code.code_strings
            if cs.file_path != target_relative
        ]
        if target_blocks:
            rw_code = CodeStringsMarkdown(
                code_strings=target_blocks + other_blocks,
            )
    except ValueError:
        pass

    # Token-limit check for read-writable context
    rw_markdown = rw_code.markdown
    rw_tokens = encoded_tokens_len(rw_markdown)
    if rw_tokens > optim_token_limit:
        raise ValueError(READ_WRITABLE_LIMIT_ERROR)

    # Preexisting objects for code replacer
    preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]] = set(
        chain(
            *(
                find_preexisting_objects(cs.code)
                for cs in rw_code.code_strings
            ),
            *(
                find_preexisting_objects(cs.code)
                for cs in all_ctx.read_only.code_strings
            ),
        ),
    )

    # Progressive fallback for read-only context
    read_only_markdown = all_ctx.read_only.markdown
    ro_tokens = encoded_tokens_len(read_only_markdown)
    if rw_tokens + ro_tokens > optim_token_limit:
        log.debug(
            "Code context has exceeded token limit, "
            "removing docstrings from read-only code",
        )
        read_only_markdown = re_extract_from_cache(
            all_ctx.file_caches,
            CodeContextType.READ_ONLY,
            project_root,
        ).markdown
        if (
            rw_tokens + encoded_tokens_len(read_only_markdown)
            > optim_token_limit
        ):
            log.debug(
                "Code context has exceeded token limit, "
                "removing read-only code",
            )
            read_only_markdown = ""

    # Build and apply testgen token limits
    testgen_context = build_testgen_context(
        all_ctx.testgen,
        project_root,
        function_to_optimize=function_to_optimize,
    )
    if encoded_tokens_len(testgen_context.markdown) > testgen_token_limit:
        log.debug(
            "Testgen context exceeded token limit, removing docstrings",
        )
        testgen_base_no_docs = re_extract_from_cache(
            all_ctx.file_caches,
            CodeContextType.TESTGEN,
            project_root,
        )
        testgen_context = build_testgen_context(
            testgen_base_no_docs,
            project_root,
            function_to_optimize=function_to_optimize,
        )
        if encoded_tokens_len(testgen_context.markdown) > testgen_token_limit:
            log.debug(
                "Testgen context still exceeded token limit, "
                "removing enrichment",
            )
            testgen_context = build_testgen_context(
                testgen_base_no_docs,
                project_root,
                include_enrichment=False,
            )
            if (
                encoded_tokens_len(testgen_context.markdown)
                > testgen_token_limit
            ):
                raise ValueError(TESTGEN_LIMIT_ERROR)

    # Hashing
    hashing_markdown = all_ctx.hashing.markdown
    hashing_hash = hashlib.sha256(
        hashing_markdown.encode("utf-8"),
    ).hexdigest()

    all_helper_fqns = sorted(
        {fs.fully_qualified_name for fs in fto_source_list + hoh_source_list},
    )

    return CodeOptimizationContext(
        read_writable=rw_markdown,
        read_only=read_only_markdown,
        hashing=hashing_markdown,
        testgen=testgen_context.markdown,
        hashing_hash=hashing_hash,
        read_writable_code=rw_code,
        testgen_context=testgen_context,
        helper_functions=fto_source_list,
        testgen_helper_fqns=all_helper_fqns,
        preexisting_objects=preexisting_objects,
    )
