"""Candidate generation strategies for the optimization pipeline.

Contains standalone functions for requesting optimization candidates
from the AI service, including line-profiler-guided, refinement,
repair, and adaptive strategies.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

from codeflash_core import (
    Candidate,
    OptimizationRequest,
)
from codeflash_core import (
    __version__ as _core_version,
)

from .._constants import LANGUAGE_VERSION

if TYPE_CHECKING:
    from codeflash_core import AIClient

    from ..context.models import CodeOptimizationContext
    from ..testing.models import TestFiles
    from ..verification.models import OriginalCodeBaseline
    from ._context import OptimizationContext
    from ._optimizer import FunctionInput

log = logging.getLogger(__name__)


def generate_candidates(  # noqa: PLR0913
    ctx: OptimizationContext,
    function_trace_id: str,
    fn_input: FunctionInput,
    code_context: CodeOptimizationContext,
    baseline: OriginalCodeBaseline | None = None,
    test_input_examples: str | None = None,
    *,
    is_numerical: bool = False,
) -> list[Candidate]:
    """Request optimization candidates from the AI service."""
    from ..context.models import (  # noqa: PLC0415
        CodeStringsMarkdown,
    )

    baseline_runtime_ns: int | None = None
    loop_count: int | None = None
    if baseline is not None:
        baseline_runtime_ns = int(baseline.runtime)
        loop_count = baseline.benchmarking_test_results.number_of_loops()

    request = OptimizationRequest(
        source_code=code_context.read_writable_code.markdown,
        language=ctx.plugin.language_id,
        language_version=LANGUAGE_VERSION,
        context_code=code_context.read_only,
        is_async=fn_input.function.is_async,
        is_numerical_code=is_numerical,
        codeflash_version=_core_version,
        baseline_runtime_ns=baseline_runtime_ns,
        loop_count=loop_count,
        test_input_examples=test_input_examples,
    )
    try:
        raw = ctx.ai_client.get_candidates(
            request,
            trace_id=function_trace_id,
        )
    except Exception:
        log.exception(
            "AI service error for %s",
            fn_input.function.qualified_name,
        )
        return []

    # The AI service returns markdown-fenced code blocks.
    # Parse them into plain Python before replacement.
    candidates: list[Candidate] = []
    for c in raw:
        parsed = CodeStringsMarkdown.parse_markdown_code(c.code)
        if not parsed.code_strings:
            log.debug(
                "Candidate %s has no parseable code blocks",
                c.candidate_id,
            )
            continue
        plain_code = "\n\n".join(cs.code for cs in parsed.code_strings)
        candidates.append(
            Candidate(
                code=plain_code,
                explanation=c.explanation,
                candidate_id=c.candidate_id,
            ),
        )
    return candidates


async def generate_lp_candidates(  # noqa: C901, PLR0913
    ctx: OptimizationContext,
    function_trace_id: str,
    test_files: TestFiles | None,
    fn_input: FunctionInput,
    code_context: CodeOptimizationContext,
    baseline: OriginalCodeBaseline,
    test_env: dict[str, str],
) -> tuple[list[Candidate], str]:
    """Generate candidates guided by line profiler data.

    Adds ``@codeflash_line_profile`` decorators to the target
    function and helpers, runs the test suite to produce a
    ``.lprof`` binary, parses the results into markdown, then
    calls the AI service's ``/optimize-line-profiler`` endpoint.

    Returns *(candidates, baseline_lp_markdown)*.
    """
    from pathlib import Path as _Path  # noqa: PLC0415

    from ..benchmarking._line_profiling import (  # noqa: PLC0415
        add_decorator_imports,
    )
    from ..benchmarking._parse_line_profile import (  # noqa: PLC0415
        parse_line_profile_results,
    )
    from ..context.models import CodeStringsMarkdown  # noqa: PLC0415
    from ..testing._test_runner import (  # noqa: PLC0415
        async_run_line_profile_tests,
    )
    from ._function_optimizer import (  # noqa: PLC0415
        is_numerical_code,
    )

    func = fn_input.function

    # Save original source for all affected files.
    files_to_restore: dict[_Path, str] = {
        func.file_path: func.file_path.read_text("utf-8"),
    }
    for helper in code_context.helper_functions:
        hp = _Path(helper.file_path)
        if hp not in files_to_restore:
            files_to_restore[hp] = hp.read_text("utf-8")  # noqa: ASYNC240

    baseline_lp_markdown = ""
    try:
        lprof_path = add_decorator_imports(
            func,
            code_context.helper_functions,
        )

        if test_files is None:
            return [], ""

        await async_run_line_profile_tests(
            test_files=test_files,
            test_env=test_env,
            cwd=ctx.project_root,
            pytest_cmd=ctx.test_cfg.pytest_cmd,
        )

        if not lprof_path.exists():
            log.debug(
                "No .lprof file produced for %s",
                func.qualified_name,
            )
            return [], ""

        lp_data, _ = parse_line_profile_results(lprof_path)
        lp_markdown: str = lp_data.get("str_out", "")
        if not lp_markdown:
            log.debug(
                "Empty line profiler output for %s",
                func.qualified_name,
            )
            return [], ""

        baseline_lp_markdown = lp_markdown
    except Exception:  # noqa: BLE001
        log.debug(
            "Line profiler step failed for %s",
            func.qualified_name,
            exc_info=True,
        )
        return [], ""
    finally:
        for path, original in files_to_restore.items():
            path.write_text(original, "utf-8")

    # Call the AI service with the profiler data.
    request = OptimizationRequest(
        source_code=code_context.read_writable_code.markdown,
        language=ctx.plugin.language_id,
        language_version=LANGUAGE_VERSION,
        context_code=code_context.read_only,
        is_numerical_code=is_numerical_code(
            fn_input.source_code,
            func.qualified_name,
        ),
        codeflash_version=_core_version,
    )
    try:
        raw = ctx.ai_client.optimize_with_line_profiler(
            request,
            line_profiler_results=lp_markdown,
            trace_id=function_trace_id,
        )
    except Exception:  # noqa: BLE001
        log.debug(
            "AI line-profiler optimization failed for %s",
            func.qualified_name,
            exc_info=True,
        )
        return [], baseline_lp_markdown

    candidates: list[Candidate] = []
    for c in raw:
        parsed = CodeStringsMarkdown.parse_markdown_code(c.code)
        if not parsed.code_strings:
            continue
        plain_code = "\n\n".join(cs.code for cs in parsed.code_strings)
        candidates.append(
            Candidate(
                code=plain_code,
                explanation=c.explanation,
                candidate_id=c.candidate_id,
            ),
        )
    log.info(
        "Generated %d line-profiler candidates for %s",
        len(candidates),
        func.qualified_name,
    )
    return candidates, baseline_lp_markdown


def generate_refinement_candidates(  # noqa: PLR0913
    ai_client: AIClient,
    function_trace_id: str,
    baseline_lp_markdown: str,
    valid: list[Candidate],
    eval_ctx: EvaluationContext,
    fn_input: FunctionInput,
    baseline: OriginalCodeBaseline,
    code_context: CodeOptimizationContext,
) -> list[Candidate]:
    """Request refined versions of valid candidates from the AI."""
    if not valid:
        return []

    from ..ai._refinement import (  # noqa: PLC0415
        RefinementRequest,
        optimize_code_refinement,
    )

    requests: list[RefinementRequest] = []
    for candidate in valid:
        cid = candidate.candidate_id
        runtime = eval_ctx.optimized_runtimes.get(cid)
        speedup = eval_ctx.speedup_ratios.get(cid)
        if runtime is None or speedup is None:
            continue
        pct = f"{int(speedup * 100)}%"
        requests.append(
            RefinementRequest(
                optimization_id=cid,
                original_source_code=fn_input.source_code,
                read_only_dependency_code=code_context.read_only,
                original_code_runtime=baseline.runtime,
                optimized_source_code=candidate.code,
                optimized_explanation=candidate.explanation,
                optimized_code_runtime=int(runtime),
                speedup=pct,
                trace_id=function_trace_id,
                original_line_profiler_results=(baseline_lp_markdown),
                optimized_line_profiler_results="",
            ),
        )
    if not requests:
        return []

    try:
        refined = optimize_code_refinement(
            ai_client,
            requests,
        )
    except Exception:  # noqa: BLE001
        log.debug(
            "Refinement failed for %s",
            fn_input.function.qualified_name,
            exc_info=True,
        )
        return []
    log.info(
        "Generated %d refinement candidates for %s",
        len(refined),
        fn_input.function.qualified_name,
    )
    return refined


def repair_failed_candidates(
    ai_client: AIClient,
    function_trace_id: str,
    failed_candidate_diffs: dict[str, list[Any]],
    failed_candidate_code: dict[str, str],
    fn_input: FunctionInput,
) -> list[Candidate]:
    """Attempt to repair candidates that failed behavioral tests."""
    if not failed_candidate_diffs:
        return []

    from ..ai._refinement import (  # noqa: PLC0415
        CodeRepairRequest,
        code_repair,
    )

    repaired: list[Candidate] = []
    for cid, diffs in failed_candidate_diffs.items():
        candidate_code = failed_candidate_code.get(cid)
        if not candidate_code:
            continue
        request = CodeRepairRequest(
            optimization_id=cid,
            original_source_code=fn_input.source_code,
            modified_source_code=candidate_code,
            trace_id=function_trace_id,
            test_diffs=tuple(diffs),
        )
        try:
            result = code_repair(ai_client, request)
        except Exception:  # noqa: BLE001
            log.warning(
                "Repair failed for candidate %s",
                cid,
                exc_info=True,
            )
            continue
        if result is not None:
            repaired.append(result)

    log.info(
        "Repaired %d candidates for %s",
        len(repaired),
        fn_input.function.qualified_name,
    )
    return repaired


def generate_adaptive_candidate(
    ai_client: AIClient,
    function_trace_id: str,
    valid: list[Candidate],
    eval_ctx: EvaluationContext,
    fn_input: FunctionInput,
) -> list[Candidate]:
    """Synthesize a new candidate from multiple valid ones."""
    if len(valid) < 2:  # noqa: PLR2004
        return []

    from ..ai._refinement import (  # noqa: PLC0415
        AdaptiveCandidate,
        AdaptiveOptimizeRequest,
        OptimizedCandidateSource,
        adaptive_optimize,
    )

    adaptive_candidates: list[AdaptiveCandidate] = []
    for candidate in valid:
        cid = candidate.candidate_id
        speedup = eval_ctx.speedup_ratios.get(cid)
        pct = f"{int(speedup * 100)}%" if speedup else "0%"
        try:
            source = OptimizedCandidateSource(
                candidate.source or OptimizedCandidateSource.OPTIMIZE.value,
            )
        except ValueError:
            source = OptimizedCandidateSource.OPTIMIZE
        adaptive_candidates.append(
            AdaptiveCandidate(
                optimization_id=cid,
                source_code=candidate.code,
                explanation=candidate.explanation,
                source=source,
                speedup=pct,
            ),
        )

    request = AdaptiveOptimizeRequest(
        trace_id=function_trace_id,
        original_source_code=fn_input.source_code,
        candidates=tuple(adaptive_candidates),
    )
    try:
        result = adaptive_optimize(ai_client, request)
    except Exception:  # noqa: BLE001
        log.debug(
            "Adaptive optimization failed for %s",
            fn_input.function.qualified_name,
            exc_info=True,
        )
        return []
    if result is None:
        return []
    log.info(
        "Generated adaptive candidate for %s",
        fn_input.function.qualified_name,
    )
    return [result]


if TYPE_CHECKING:
    from typing import Any

    from codeflash_core import EvaluationContext
