"""
Module: sandbox.backends.llm
Description: Local LLM backend that calls provider APIs directly via httpx.

Implements:
    - docs/cli/reference.md#huitzo-dev
    - docs/guides/developer-environment.md#development-session-flow

See Also:
    - docs/sdk/integrations.md#llm-integration
"""

from __future__ import annotations

import asyncio
import os
from collections.abc import AsyncIterator
from typing import Any

import httpx

from huitzo_sdk.errors import LLMError

_OPENAI_URL = "https://api.openai.com/v1/chat/completions"
_ANTHROPIC_URL = "https://api.anthropic.com/v1/messages"

_OPENAI_PREFIXES = ("gpt-", "o1-", "o3-", "o4-")
_ANTHROPIC_PREFIXES = ("claude-",)


def _detect_provider(model: str) -> str:
    """Detect LLM provider from model name."""
    lower = model.lower()
    if any(lower.startswith(p) for p in _OPENAI_PREFIXES):
        return "openai"
    if any(lower.startswith(p) for p in _ANTHROPIC_PREFIXES):
        return "anthropic"
    raise LLMError(
        provider="unknown",
        model=model,
        message=f"Cannot detect provider for model '{model}'. "
        "Expected prefix: gpt-*, o1-*, o3-*, o4-* (OpenAI) or claude-* (Anthropic).",
    )


def _get_api_key(provider: str) -> str:
    """Read API key from environment."""
    env_var = "OPENAI_API_KEY" if provider == "openai" else "ANTHROPIC_API_KEY"
    key = os.environ.get(env_var)
    if not key:
        raise LLMError(
            provider=provider,
            message=f"Set {env_var} environment variable to use {provider} models in sandbox.",
        )
    return key


class LocalLLMBackend:
    """Local LLM backend that calls OpenAI/Anthropic APIs directly.

    Satisfies the huitzo_sdk LLMBackend protocol.
    """

    def __init__(self) -> None:
        self._client: httpx.AsyncClient | None = None
        self._lock = asyncio.Lock()

    async def _get_client(self) -> httpx.AsyncClient:
        if self._client is None:
            async with self._lock:
                if self._client is None:
                    self._client = httpx.AsyncClient(timeout=120.0)
        return self._client

    async def complete(
        self,
        *,
        prompt: str,
        system: str | None,
        model: str,
        temperature: float,
        max_tokens: int | None,
        top_p: float | None,
        stop: list[str] | None,
        response_format: str | None,
        schema: type[Any] | None,
    ) -> str:
        provider = _detect_provider(model)
        api_key = _get_api_key(provider)
        client = await self._get_client()

        if provider == "openai":
            return await self._complete_openai(
                client,
                api_key,
                prompt=prompt,
                system=system,
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                top_p=top_p,
                stop=stop,
            )
        return await self._complete_anthropic(
            client,
            api_key,
            prompt=prompt,
            system=system,
            model=model,
            temperature=temperature,
            max_tokens=max_tokens,
            top_p=top_p,
            stop=stop,
        )

    async def _complete_openai(
        self,
        client: httpx.AsyncClient,
        api_key: str,
        *,
        prompt: str,
        system: str | None,
        model: str,
        temperature: float,
        max_tokens: int | None,
        top_p: float | None,
        stop: list[str] | None,
    ) -> str:
        messages: list[dict[str, str]] = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        payload: dict[str, Any] = {"model": model, "messages": messages, "temperature": temperature}
        if max_tokens is not None:
            payload["max_tokens"] = max_tokens
        if top_p is not None:
            payload["top_p"] = top_p
        if stop:
            payload["stop"] = stop

        resp = await client.post(
            _OPENAI_URL,
            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
            json=payload,
        )
        if resp.status_code != 200:
            raise LLMError(
                provider="openai",
                model=model,
                status_code=resp.status_code,
                message=f"OpenAI API error: {resp.text[:200]}",
            )
        data = resp.json()
        result: str = data["choices"][0]["message"]["content"]
        return result

    async def _complete_anthropic(
        self,
        client: httpx.AsyncClient,
        api_key: str,
        *,
        prompt: str,
        system: str | None,
        model: str,
        temperature: float,
        max_tokens: int | None,
        top_p: float | None,
        stop: list[str] | None,
    ) -> str:
        payload: dict[str, Any] = {
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens or 4096,
            "temperature": temperature,
        }
        if system:
            payload["system"] = system
        if top_p is not None:
            payload["top_p"] = top_p
        if stop:
            payload["stop_sequences"] = stop

        resp = await client.post(
            _ANTHROPIC_URL,
            headers={
                "x-api-key": api_key,
                "anthropic-version": "2023-06-01",
                "Content-Type": "application/json",
            },
            json=payload,
        )
        if resp.status_code != 200:
            raise LLMError(
                provider="anthropic",
                model=model,
                status_code=resp.status_code,
                message=f"Anthropic API error: {resp.text[:200]}",
            )
        data = resp.json()
        result: str = data["content"][0]["text"]
        return result

    async def stream(
        self,
        *,
        prompt: str,
        system: str | None,
        model: str,
        temperature: float,
        max_tokens: int | None,
        top_p: float | None,
        stop: list[str] | None,
    ) -> AsyncIterator[str]:
        provider = _detect_provider(model)
        api_key = _get_api_key(provider)
        client = await self._get_client()

        if provider == "openai":
            async for chunk in self._stream_openai(
                client,
                api_key,
                prompt=prompt,
                system=system,
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                top_p=top_p,
                stop=stop,
            ):
                yield chunk
        else:
            async for chunk in self._stream_anthropic(
                client,
                api_key,
                prompt=prompt,
                system=system,
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                top_p=top_p,
                stop=stop,
            ):
                yield chunk

    async def _stream_openai(
        self,
        client: httpx.AsyncClient,
        api_key: str,
        *,
        prompt: str,
        system: str | None,
        model: str,
        temperature: float,
        max_tokens: int | None,
        top_p: float | None,
        stop: list[str] | None,
    ) -> AsyncIterator[str]:
        messages: list[dict[str, str]] = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "stream": True,
        }
        if max_tokens is not None:
            payload["max_tokens"] = max_tokens
        if top_p is not None:
            payload["top_p"] = top_p
        if stop:
            payload["stop"] = stop

        async with client.stream(
            "POST",
            _OPENAI_URL,
            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
            json=payload,
        ) as resp:
            if resp.status_code != 200:
                body = await resp.aread()
                raise LLMError(
                    provider="openai",
                    model=model,
                    status_code=resp.status_code,
                    message=f"OpenAI API error: {body.decode()[:200]}",
                )
            async for line in resp.aiter_lines():
                if line.startswith("data: ") and line != "data: [DONE]":
                    import json

                    chunk = json.loads(line[6:])
                    delta = chunk.get("choices", [{}])[0].get("delta", {})
                    content = delta.get("content")
                    if content:
                        yield content

    async def _stream_anthropic(
        self,
        client: httpx.AsyncClient,
        api_key: str,
        *,
        prompt: str,
        system: str | None,
        model: str,
        temperature: float,
        max_tokens: int | None,
        top_p: float | None,
        stop: list[str] | None,
    ) -> AsyncIterator[str]:
        payload: dict[str, Any] = {
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens or 4096,
            "temperature": temperature,
            "stream": True,
        }
        if system:
            payload["system"] = system
        if top_p is not None:
            payload["top_p"] = top_p
        if stop:
            payload["stop_sequences"] = stop

        async with client.stream(
            "POST",
            _ANTHROPIC_URL,
            headers={
                "x-api-key": api_key,
                "anthropic-version": "2023-06-01",
                "Content-Type": "application/json",
            },
            json=payload,
        ) as resp:
            if resp.status_code != 200:
                body = await resp.aread()
                raise LLMError(
                    provider="anthropic",
                    model=model,
                    status_code=resp.status_code,
                    message=f"Anthropic API error: {body.decode()[:200]}",
                )
            async for line in resp.aiter_lines():
                if line.startswith("data: "):
                    import json

                    event = json.loads(line[6:])
                    if event.get("type") == "content_block_delta":
                        text = event.get("delta", {}).get("text")
                        if text:
                            yield text
