# SPDX-FileCopyrightText: 2025 Weibo, Inc.
#
# SPDX-License-Identifier: Apache-2.0

"""
Base LLM Provider interface for Simple Chat.

Defines the abstract interface for all LLM providers with streaming support.
This is a simplified version without tool calling support.
"""

import asyncio
import json
import logging
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, AsyncGenerator

import httpx

logger = logging.getLogger(__name__)


class ChunkType(Enum):
    """Type of streaming chunk."""

    CONTENT = "content"
    ERROR = "error"
    DONE = "done"


@dataclass
class StreamChunk:
    """Unified streaming chunk representation."""

    type: ChunkType
    content: str = ""
    error: str | None = None


@dataclass
class ProviderConfig:
    """Configuration for LLM provider."""

    api_key: str
    base_url: str
    model_id: str
    default_headers: dict[str, Any] = field(default_factory=dict)
    timeout: float = 300.0
    max_tokens: int = 32768


class LLMProvider(ABC):
    """
    Abstract base class for LLM providers.

    Defines the interface for streaming chat completions.
    This is a simplified version without tool calling support.
    """

    def __init__(self, config: ProviderConfig, client: httpx.AsyncClient):
        self.config = config
        self.client = client

    @property
    @abstractmethod
    def provider_name(self) -> str:
        """Return the provider name (e.g., 'openai', 'claude', 'gemini')."""
        pass

    @abstractmethod
    async def stream_chat(
        self,
        messages: list[dict[str, Any]],
        cancel_event: asyncio.Event,
    ) -> AsyncGenerator[StreamChunk, None]:
        """Stream chat completion from the LLM."""
        pass

    @abstractmethod
    def format_messages(self, messages: list[dict[str, Any]]) -> Any:
        """Format messages for this provider's API."""
        pass

    def _build_headers(self) -> dict[str, str]:
        """Build HTTP headers for API requests."""
        headers = {"Content-Type": "application/json"}
        if self.config.default_headers:
            headers.update(self.config.default_headers)
        return headers

    async def _check_cancellation(self, cancel_event: asyncio.Event) -> bool:
        """Check if cancellation has been requested."""
        return cancel_event.is_set()

    async def _stream_sse(
        self,
        url: str,
        payload: dict[str, Any],
        headers: dict[str, str],
        cancel_event: asyncio.Event,
    ) -> AsyncGenerator[dict[str, Any], None]:
        """
        Stream SSE responses from API.

        Common SSE parsing logic shared by all providers.

        Args:
            url: API endpoint URL
            payload: Request payload
            headers: Request headers
            cancel_event: Cancellation event

        Yields:
            Parsed JSON data from SSE events
        """
        start_time = time.time()
        chunk_count = 0

        try:
            async with self.client.stream(
                "POST", url, json=payload, headers=headers
            ) as response:
                if response.status_code >= 400:
                    error_body = await response.aread()
                    error_msg = error_body.decode("utf-8", errors="replace")
                    logger.error(
                        "%s API error: status=%s, body=%s",
                        self.provider_name,
                        response.status_code,
                        error_msg,
                    )
                    yield {"_error": error_msg}
                    return

                response.raise_for_status()

                async for line in response.aiter_lines():
                    if await self._check_cancellation(cancel_event):
                        logger.info(
                            "%s stream cancelled after %d chunks in %.2fs",
                            self.provider_name,
                            chunk_count,
                            time.time() - start_time,
                        )
                        return

                    if not line or line.startswith(":"):
                        continue

                    # Handle both "data: " (with space) and "data:" (without space)
                    if line.startswith("data:"):
                        # Strip "data:" prefix and any leading whitespace
                        data = line[5:].lstrip()
                        if data == "[DONE]":
                            logger.info(
                                "%s stream completed: %d chunks in %.2fs",
                                self.provider_name,
                                chunk_count,
                                time.time() - start_time,
                            )
                            return

                        try:
                            chunk_count += 1
                            yield json.loads(data)
                        except json.JSONDecodeError:
                            continue

                # If we exit the loop without [DONE], log it
                logger.info(
                    "%s stream ended without [DONE] marker after %d chunks in %.2fs",
                    self.provider_name,
                    chunk_count,
                    time.time() - start_time,
                )

        except httpx.RequestError as e:
            logger.exception("%s request error", self.provider_name)
            yield {"_error": str(e)}
