from __future__ import annotations

import json
import logging
from typing import Any, Dict, List, Optional

from actor.observability import ComponentStatus, ObservabilityService

logger = logging.getLogger(__name__)

TOKEN_LIMIT_CASCADE = [200_000, 128_000, 96_000, 64_000, 48_000, 32_000, 24_000, 16_000, 8_000, 4_096]


class LLMClient:
    """
    Simple OpenAI-compatible LLM client for structured and unstructured chat.
    """

    def __init__(
        self,
        endpoint: str,
        auth_type: str,
        model: str,
        observability: ObservabilityService,
        api_key: Optional[str] = None,
        username: Optional[str] = None,
        password: Optional[str] = None,
        api_version: Optional[str] = None,
        deployment_name: Optional[str] = None,
    ):
        self.endpoint = endpoint
        self.auth_type = auth_type
        self.api_key = api_key
        self.username = username
        self.password = password
        self.api_version = api_version
        self.deployment_name = deployment_name
        self.model = model
        self.observability = observability

        from actor.connectors._openai_common import build_openai_client
        client, effective_model, _ = build_openai_client(
            endpoint=endpoint,
            auth_type=auth_type,
            api_key=api_key,
            username=username,
            password=password,
            api_version=api_version,
            deployment_name=deployment_name,
            timeout=60.0,
        )
        self.client = client
        if effective_model:
            self.model = effective_model
        self.observability.register_component("llm")
        logger.info("LLM client initialized", extra={"endpoint": endpoint, "model": model})

    def completion(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                **kwargs,
            )
            return {
                "content": response.choices[0].message.content,
                "usage": {
                    "prompt_tokens": response.usage.prompt_tokens,
                    "completion_tokens": response.usage.completion_tokens,
                    "total_tokens": response.usage.total_tokens,
                },
            }
        except Exception as exc:  # pragma: no cover - live dependency
            logger.error("LLM completion failed: %s", exc)
            self.observability.update_component_health(
                component="llm",
                status=ComponentStatus.UNHEALTHY,
                details={"last_error": str(exc)},
                error=str(exc),
            )
            raise

    def structured_completion(
        self,
        messages: List[Dict[str, str]],
        json_schema: Dict[str, Any],
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                response_format={"type": "json_schema", "json_schema": json_schema},
                **kwargs,
            )
            content = response.choices[0].message.content
            parsed = json.loads(content)
            return {
                "content": content,
                "parsed": parsed,
                "usage": {
                    "prompt_tokens": response.usage.prompt_tokens,
                    "completion_tokens": response.usage.completion_tokens,
                    "total_tokens": response.usage.total_tokens,
                },
            }
        except Exception as exc:  # pragma: no cover
            logger.error("LLM structured completion failed: %s", exc)
            self.observability.update_component_health(
                component="llm",
                status=ComponentStatus.UNHEALTHY,
                details={"last_error": str(exc)},
                error=str(exc),
            )
            raise

    def validate(self) -> None:
        import time

        start = time.time()
        try:
            resp = self.client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": "test"}],
                max_tokens=5,
            )
            latency_ms = int((time.time() - start) * 1000)
            status = ComponentStatus.HEALTHY if latency_ms <= 5000 else ComponentStatus.DEGRADED
            self.observability.update_component_health(
                component="llm",
                status=status,
                details={
                    "endpoint": self.endpoint,
                    "model": self.model,
                    "connected": True,
                    "test_latency_ms": latency_ms,
                },
            )
        except Exception as exc:  # pragma: no cover
            self.observability.update_component_health(
                component="llm",
                status=ComponentStatus.UNHEALTHY,
                details={"endpoint": self.endpoint, "model": self.model, "connected": False},
                error=str(exc),
            )
            raise

    @staticmethod
    def scale_down_token_limit(current: int) -> int:
        """Find next lower limit in cascade when token exceeded."""
        for limit in TOKEN_LIMIT_CASCADE:
            if limit < current:
                return limit
        return 4096
