"""HTTP client for communicating with the Codeflash AI optimization service."""

from __future__ import annotations

import contextlib
import sys
import uuid
from typing import TYPE_CHECKING, Any

if sys.version_info >= (3, 11):
    from typing import Self
else:
    from typing_extensions import Self

import attrs

from ._http import (
    _make_session,
    _resolve_api_key,
    _resolve_base_url,
    _strip_trailing_slash,
)
from ._model import Candidate, OptimizationRequest, OptimizationReviewResult
from .exceptions import (
    AIServiceConnectionError,
    AIServiceError,
)


def _parse_candidates(data: dict[str, Any]) -> list[Candidate]:
    """Build a list of candidates from an AI service response."""
    return [
        Candidate(
            code=item.get("source_code", ""),
            explanation=item.get("explanation", ""),
            candidate_id=item.get("optimization_id", ""),
        )
        for item in data.get("optimizations", [])
    ]


if TYPE_CHECKING:
    import requests as requests_types


@attrs.define
class AIClient:
    """
    HTTP client for the Codeflash AI optimization service.
    """

    _base_url: str = attrs.field(
        alias="base_url",
        default=attrs.Factory(_resolve_base_url),
        converter=_strip_trailing_slash,
    )
    _api_key: str = attrs.field(
        alias="api_key",
        default=attrs.Factory(_resolve_api_key),
        repr=False,
    )
    _timeout: float = attrs.field(
        alias="timeout",
        default=300.0,
    )
    _session: requests_types.Session = attrs.field(
        init=False,
        factory=_make_session,
    )

    def __attrs_post_init__(self) -> None:
        """Set the Authorization header on the session."""
        if self._api_key:
            self._session.headers["Authorization"] = f"Bearer {self._api_key}"

    def __enter__(self) -> Self:
        """Enter the context manager."""
        return self

    def __exit__(self, *exc_info: object) -> None:
        """Exit the context manager and close the session."""
        self.close()

    def post(
        self,
        endpoint: str,
        payload: dict[str, Any] | list[Any],
    ) -> dict[str, Any]:
        """POST to ``/ai{endpoint}`` and return the JSON response.

        Centralizes HTTP call + error handling for all AI
        service endpoints.  *endpoint* should start with ``/``
        (e.g. ``/optimize``, ``/rank``).
        """
        import requests  # noqa: PLC0415

        url = f"{self._base_url}/ai{endpoint}"
        try:
            resp = self._session.post(
                url,
                json=payload,
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return resp.json()  # type: ignore[no-any-return]
        except requests.HTTPError:
            raise AIServiceError(
                resp.status_code,
                resp.text,
            ) from None
        except requests.RequestException as e:
            raise AIServiceConnectionError(str(e)) from e
        except ValueError:
            raise AIServiceError(
                resp.status_code,
                resp.text,
            ) from None

    def get_candidates(
        self,
        request: OptimizationRequest,
        n_candidates: int = 5,
        trace_id: str = "",
    ) -> list[Candidate]:
        """
        Request optimization candidates for *request*.
        """
        if not trace_id:
            trace_id = str(uuid.uuid4())

        payload: dict[str, Any] = {
            "source_code": request.source_code,
            "dependency_code": request.context_code,
            "trace_id": trace_id,
            "language": request.language,
            "language_version": request.language_version,
            "n_candidates": n_candidates,
            "call_sequence": 1,
            "is_async": request.is_async,
            "is_numerical_code": request.is_numerical_code,
            "codeflash_version": request.codeflash_version,
        }
        if request.baseline_runtime_ns is not None:
            payload["baseline_runtime_ns"] = request.baseline_runtime_ns
        if request.loop_count is not None:
            payload["loop_count"] = request.loop_count
        if request.line_profiler_results is not None:
            payload["line_profiler_results"] = request.line_profiler_results
        if request.test_input_examples is not None:
            payload["test_input_examples"] = request.test_input_examples
        data = self.post("/optimize", payload)
        return _parse_candidates(data)

    def generate_ranking(
        self,
        trace_id: str,
        diffs: list[str],
        candidate_ids: list[str],
        speedups: list[float],
    ) -> list[int] | None:
        """Rank optimization candidates via the AI service.

        Returns candidate indices in decreasing quality order,
        or *None* if the service is unavailable.
        """
        payload = {
            "trace_id": trace_id,
            "diffs": diffs,
            "speedups": speedups,
            "optimization_ids": candidate_ids,
        }
        try:
            data = self.post("/rank", payload)
        except (AIServiceError, AIServiceConnectionError):
            return None

        ranking: list[int] = data.get("ranking", [])
        return ranking or None

    def optimize_with_line_profiler(
        self,
        request: OptimizationRequest,
        line_profiler_results: str,
        n_candidates: int = 5,
        trace_id: str = "",
    ) -> list[Candidate]:
        """
        Request optimization candidates guided by *line_profiler_results*.
        """
        if not line_profiler_results:
            return []

        if not trace_id:
            trace_id = str(uuid.uuid4())

        payload = {
            "source_code": request.source_code,
            "dependency_code": request.context_code,
            "n_candidates": n_candidates,
            "line_profiler_results": line_profiler_results,
            "trace_id": trace_id,
            "language": request.language,
            "language_version": request.language_version,
            "call_sequence": 1,
            "is_numerical_code": request.is_numerical_code,
            "codeflash_version": request.codeflash_version,
        }
        data = self.post("/optimize-line-profiler", payload)
        return _parse_candidates(data)

    def generate_explanation(
        self,
        payload: dict[str, Any],
    ) -> str:
        """
        Request an updated explanation from the AI service.

        Returns the explanation text, or ``""`` on failure.
        """
        try:
            data = self.post("/explain", payload)
        except (AIServiceError, AIServiceConnectionError):
            return ""
        explanation: str = data.get("explanation", "")
        return explanation

    def log_results(
        self,
        payload: dict[str, Any],
    ) -> None:
        """
        Log optimization results to the AI service (fire-and-forget).
        """
        with contextlib.suppress(AIServiceError, AIServiceConnectionError):
            self.post("/log_features", payload)

    def get_optimization_review(
        self,
        payload: dict[str, Any],
    ) -> OptimizationReviewResult:
        """
        Request an optimization quality review.

        Returns an *OptimizationReviewResult* with the review level
        and explanation, or empty strings on failure.
        """
        try:
            data = self.post("/optimization_review", payload)
        except (AIServiceError, AIServiceConnectionError):
            return OptimizationReviewResult(
                review="",
                explanation="",
            )
        return OptimizationReviewResult(
            review=data.get("review", ""),
            explanation=data.get("review_explanation", ""),
        )

    def close(self) -> None:
        """
        Close the underlying HTTP session.
        """
        self._session.close()
