"""AI service wrappers for refinement, repair, and adaptive optimization."""

from __future__ import annotations

import enum
import logging
from typing import TYPE_CHECKING, Any

import attrs

from .._constants import LANGUAGE_FIELDS

if TYPE_CHECKING:
    from codeflash_core import AIClient, Candidate

log = logging.getLogger(__name__)


class OptimizedCandidateSource(str, enum.Enum):
    """Source of an optimization candidate."""

    OPTIMIZE = "OPTIMIZE"
    OPTIMIZE_LP = "OPTIMIZE_LP"
    REFINE = "REFINE"
    REPAIR = "REPAIR"
    ADAPTIVE = "ADAPTIVE"
    JIT_REWRITE = "JIT_REWRITE"


@attrs.frozen
class RefinementRequest:
    """Request for code refinement."""

    optimization_id: str
    original_source_code: str
    read_only_dependency_code: str
    original_code_runtime: int
    optimized_source_code: str
    optimized_explanation: str
    optimized_code_runtime: int
    speedup: str
    trace_id: str
    original_line_profiler_results: str
    optimized_line_profiler_results: str
    function_references: str | None = None
    additional_context_files: dict[str, str] | None = None


@attrs.frozen
class CodeRepairRequest:
    """Request for code repair."""

    optimization_id: str
    original_source_code: str
    modified_source_code: str
    trace_id: str
    test_diffs: tuple[dict[str, object], ...]


@attrs.frozen
class AdaptiveCandidate:
    """A candidate for adaptive optimization."""

    optimization_id: str
    source_code: str
    explanation: str
    source: OptimizedCandidateSource
    speedup: str


@attrs.frozen
class AdaptiveOptimizeRequest:
    """Request for adaptive optimization."""

    trace_id: str
    original_source_code: str
    candidates: tuple[AdaptiveCandidate, ...]


def _parse_candidate(data: dict[str, Any], source: str) -> Candidate | None:
    """Parse a response dict into a Candidate.

    Runs markdown code-block parsing on the raw ``source_code``
    field, matching the pattern used by ``_candidate_gen.py``.
    """
    from codeflash_core import Candidate  # noqa: PLC0415

    from ..context.models import CodeStringsMarkdown  # noqa: PLC0415

    raw_code = data.get("source_code", "")
    if not raw_code:
        return None

    parsed = CodeStringsMarkdown.parse_markdown_code(raw_code)
    if parsed.code_strings:
        code = "\n\n".join(cs.code for cs in parsed.code_strings)
    else:
        code = raw_code

    return Candidate(
        code=code,
        explanation=data.get("explanation", ""),
        candidate_id=data.get("optimization_id", ""),
        source=source,
    )


def optimize_code_refinement(
    client: AIClient,
    requests_: list[RefinementRequest],
    rerun_trace_id: str | None = None,
) -> list[Candidate]:
    """
    Send refinement requests to the AI service.

    Returns a list of refined candidates.
    """
    from codeflash_core import humanize_runtime  # noqa: PLC0415

    payload: list[dict[str, Any]] = []
    for req in requests_:
        item: dict[str, Any] = {
            **LANGUAGE_FIELDS,
            "optimization_id": req.optimization_id,
            "original_source_code": req.original_source_code,
            "read_only_dependency_code": req.read_only_dependency_code,
            "original_line_profiler_results": (
                req.original_line_profiler_results
            ),
            "original_code_runtime": humanize_runtime(
                req.original_code_runtime,
            ),
            "optimized_source_code": req.optimized_source_code,
            "optimized_explanation": req.optimized_explanation,
            "optimized_line_profiler_results": (
                req.optimized_line_profiler_results
            ),
            "optimized_code_runtime": humanize_runtime(
                req.optimized_code_runtime,
            ),
            "speedup": req.speedup,
            "trace_id": req.trace_id,
            "function_references": req.function_references,
            "call_sequence": 1,
            "rerun_trace_id": rerun_trace_id,
        }
        if req.additional_context_files:
            item["additional_context_files"] = req.additional_context_files
        payload.append(item)

    data = client.post("/refinement", payload)

    candidates: list[Candidate] = []
    for item in data.get("refinements", []):
        candidate = _parse_candidate(
            item, OptimizedCandidateSource.REFINE.value
        )
        if candidate is not None:
            candidates.append(candidate)
    return candidates


def code_repair(
    client: AIClient,
    request: CodeRepairRequest,
) -> Candidate | None:
    """
    Send a repair request to the AI service.

    Returns a repaired candidate or *None*.
    """
    payload: dict[str, Any] = {
        **LANGUAGE_FIELDS,
        "optimization_id": request.optimization_id,
        "original_source_code": request.original_source_code,
        "modified_source_code": request.modified_source_code,
        "trace_id": request.trace_id,
        "test_diffs": list(request.test_diffs),
    }

    data = client.post("/code_repair", payload)

    return _parse_candidate(data, OptimizedCandidateSource.REPAIR.value)


def adaptive_optimize(
    client: AIClient,
    request: AdaptiveOptimizeRequest,
) -> Candidate | None:
    """
    Send an adaptive optimization request to the AI service.

    Returns a new candidate or *None*.
    """
    payload: dict[str, Any] = {
        "trace_id": request.trace_id,
        "original_source_code": request.original_source_code,
        "candidates": [attrs.asdict(c) for c in request.candidates],
    }

    data = client.post("/adaptive_optimize", payload)

    return _parse_candidate(data, OptimizedCandidateSource.ADAPTIVE.value)
