from __future__ import annotations

import asyncio
import logging
import os
import time
from typing import TYPE_CHECKING, Any

import attrs

from codeflash_api.llm._cost import calculate_llm_cost
from codeflash_api.llm._retry import (
    ANTHROPIC_MAX_INPUT_TOKENS,
    CHARS_PER_TOKEN_ESTIMATE,
)
from codeflash_api.observability._background import fire_and_forget
from codeflash_api.observability._recording import record_llm_call

if TYPE_CHECKING:
    from anthropic import AsyncAnthropicBedrock
    from openai import AsyncAzureOpenAI

    from codeflash_api.llm._models import LLM

log = logging.getLogger(__name__)


class LLMOutputUnparseableError(Exception):
    """
    Raised when LLM output cannot be parsed.
    """

    def __init__(
        self,
        message: str = "LLM output could not be parsed",
        *,
        cost: float = 0.0,
    ) -> None:
        super().__init__(message)
        self.cost = cost


@attrs.frozen
class LLMUsage:
    """
    Token counts from an LLM call.
    """

    input_tokens: int
    output_tokens: int


@attrs.frozen
class LLMResponse:
    """
    Response from an LLM call.
    """

    content: str
    usage: LLMUsage
    raw_response: Any
    cost: float = 0.0


@attrs.define
class LLMClient:
    """
    Async LLM client supporting OpenAI and Anthropic.
    """

    _openai_client: AsyncAzureOpenAI | None = None
    _anthropic_client: AsyncAnthropicBedrock | None = None
    _client_loop: asyncio.AbstractEventLoop | None = None
    _background_tasks: set[asyncio.Task[Any]] = attrs.Factory(set)

    def _ensure_clients(self) -> None:
        """
        Recreate clients if the event loop changed.
        """
        loop = asyncio.get_running_loop()
        if loop is self._client_loop:
            return

        self._client_loop = loop
        self._background_tasks = set()

        if os.environ.get("AZURE_OPENAI_API_KEY"):
            from openai import AsyncAzureOpenAI

            self._openai_client = AsyncAzureOpenAI()
        else:
            self._openai_client = None

        if os.environ.get("AWS_ACCESS_KEY_ID") and os.environ.get(
            "AWS_SECRET_ACCESS_KEY"
        ):
            from anthropic import AsyncAnthropicBedrock

            self._anthropic_client = AsyncAnthropicBedrock(
                aws_access_key=os.environ.get("AWS_ACCESS_KEY_ID", ""),
                aws_secret_key=os.environ.get("AWS_SECRET_ACCESS_KEY", ""),
                aws_region=os.environ.get("AWS_REGION", "us-east-1"),
            )
        else:
            self._anthropic_client = None

    async def call(
        self,
        llm: LLM,
        messages: list[dict[str, Any]],
        *,
        max_tokens: int = 16384,
        call_type: str = "",
        trace_id: str = "",
        user_id: str | None = None,
    ) -> LLMResponse:
        """
        Route to the correct provider and return an LLMResponse.
        """
        self._ensure_clients()
        start = time.monotonic()
        result: LLMResponse | None = None
        error: Exception | None = None

        try:
            if llm.model_type == "anthropic":
                result = await self._call_anthropic(llm, messages, max_tokens)
            elif llm.model_type == "openai":
                result = await self._call_openai(llm, messages, max_tokens)
            else:
                msg = f"Unsupported model type: {llm.model_type}"
                raise ValueError(msg)

            result = LLMResponse(
                content=result.content,
                usage=result.usage,
                raw_response=result.raw_response,
                cost=calculate_llm_cost(result.raw_response, llm),
            )

        except Exception as exc:
            error = exc
            latency_ms = int((time.monotonic() - start) * 1000)
            log.exception(
                "LLM call failed: type=%s trace=%s latency=%dms",
                call_type,
                trace_id,
                latency_ms,
            )
            raise

        finally:
            latency_ms = int((time.monotonic() - start) * 1000)
            fire_and_forget(
                record_llm_call(
                    trace_id=trace_id,
                    call_type=call_type,
                    model_name=llm.name,
                    messages=messages,
                    latency_ms=latency_ms,
                    user_id=user_id,
                    result=result,
                    error=error,
                    llm_cost=result.cost if result else None,
                )
            )

        return result

    async def _call_anthropic(
        self,
        llm: LLM,
        messages: list[dict[str, Any]],
        max_tokens: int,
    ) -> LLMResponse:
        """
        Call Anthropic via Bedrock.
        """
        if self._anthropic_client is None:
            msg = "Anthropic client not configured"
            raise RuntimeError(msg)

        estimated_tokens = (
            sum(len(str(m.get("content", ""))) for m in messages)
            // CHARS_PER_TOKEN_ESTIMATE
        )

        if estimated_tokens > ANTHROPIC_MAX_INPUT_TOKENS:
            msg = (
                f"Prompt too large"
                f" (~{estimated_tokens} tokens"
                f" estimated, limit"
                f" {ANTHROPIC_MAX_INPUT_TOKENS})"
            )
            raise ValueError(msg)

        system_prompt = next(
            (m["content"] for m in messages if m["role"] == "system"),
            None,
        )
        non_system = [
            {"role": m["role"], "content": m["content"]}
            for m in messages
            if m["role"] != "system"
        ]

        kwargs: dict[str, Any] = {
            "model": llm.name,
            "messages": non_system,
            "max_tokens": max_tokens,
        }
        if system_prompt:
            kwargs["system"] = system_prompt

        response = await self._anthropic_client.messages.create(**kwargs)
        content = "".join(
            block.text for block in response.content if hasattr(block, "text")
        )
        if not content:
            cost = calculate_llm_cost(response, llm)
            raise LLMOutputUnparseableError(cost=cost)

        return LLMResponse(
            content=content,
            usage=LLMUsage(
                input_tokens=response.usage.input_tokens,
                output_tokens=response.usage.output_tokens,
            ),
            raw_response=response,
        )

    async def _call_openai(
        self,
        llm: LLM,
        messages: list[dict[str, Any]],
        max_tokens: int,
    ) -> LLMResponse:
        """
        Call Azure OpenAI.
        """
        if self._openai_client is None:
            msg = "OpenAI client not configured"
            raise RuntimeError(msg)

        if llm.name == "gpt-5-mini":
            response = await self._openai_client.chat.completions.create(
                model=llm.name,
                messages=messages,  # type: ignore[arg-type]
                max_completion_tokens=max_tokens,
            )
        else:
            response = await self._openai_client.chat.completions.create(
                model=llm.name,
                messages=messages,  # type: ignore[arg-type]
                max_tokens=max_tokens,
            )

        content = (
            response.choices[0].message.content if response.choices else None
        )
        if not content:
            cost = calculate_llm_cost(response, llm)
            raise LLMOutputUnparseableError(cost=cost)

        return LLMResponse(
            content=content,
            usage=LLMUsage(
                input_tokens=(
                    response.usage.prompt_tokens if response.usage else 0
                ),
                output_tokens=(
                    response.usage.completion_tokens if response.usage else 0
                ),
            ),
            raw_response=response,
        )
