import json
import requests
from typing import Any, Iterator, Optional, List, Union
from .errors import BeaverError
from .types import (
    ChatCompletionResponse, 
    Choice, 
    Message, 
    Usage, 
    ChatCompletionChunk, 
    StreamChoice, 
    StreamDelta
)


class ChatCompletions:
    def __init__(self, client):
        self._client = client

    def create(
        self,
        model: str,
        messages: List[Union[dict, Message]],
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        max_output_tokens: Optional[int] = None,  # Deprecated alias
        top_p: Optional[float] = None,
        stream: Optional[bool] = None,
    ) -> Union[ChatCompletionResponse, Iterator[ChatCompletionChunk]]:
        """
        Create a chat completion.
        """
        # Convert Message objects to dicts if necessary
        formatted_messages = [
            m.to_dict() if isinstance(m, Message) else m for m in messages
        ]

        request_body: dict[str, Any] = {
            "model": model,
            "messages": formatted_messages,
        }

        if temperature is not None:
            request_body["temperature"] = temperature
        
        # Standardize max_tokens
        final_max_tokens = max_tokens or max_output_tokens
        if final_max_tokens is not None:
            request_body["max_tokens"] = final_max_tokens
            
        if top_p is not None:
            request_body["top_p"] = top_p

        if stream is False:
            request_body["stream"] = False

        # Non-streaming path
        if stream is False:
            raw = self._client._post("/v1/chat/completions", request_body)
            
            # Map to ChatCompletionResponse object
            return ChatCompletionResponse(
                id=raw["id"],
                model=raw["model"],
                choices=[
                    Choice(
                        index=c["index"],
                        message=Message(
                            role=c["message"]["role"],
                            content=c["message"]["content"]
                        ),
                        finish_reason=c.get("finish_reason")
                    ) for c in raw["choices"]
                ],
                usage=Usage(**raw["usage"]) if "usage" in raw else None
            )

        # Streaming path (default)
        response = requests.post(
            f"{self._client.base_url}/v1/chat/completions",
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self._client.api_key}",
            },
            json=request_body,
            stream=True,
        )

        if not response.ok:
            self._client._handle_error(response)

        return self._parse_stream(response)

    def _parse_stream(self, response: requests.Response) -> Iterator[ChatCompletionChunk]:
        """Parse SSE stream and yield ChatCompletionChunk objects."""
        for line in response.iter_lines(decode_unicode=True):
            if not line or not line.startswith("data: "):
                continue

            data = line[6:]  # Remove "data: " prefix
            if data.strip() == "[DONE]":
                break

            try:
                chunk_raw = json.loads(data)
                
                # Map to ChatCompletionChunk
                choices = []
                for c in chunk_raw.get("choices", []):
                    delta_raw = c.get("delta", {})
                    choices.append(
                        StreamChoice(
                            index=c["index"],
                            delta=StreamDelta(
                                content=delta_raw.get("content"),
                                role=delta_raw.get("role")
                            ),
                            finish_reason=c.get("finish_reason")
                        )
                    )
                
                yield ChatCompletionChunk(
                    id=chunk_raw["id"],
                    model=chunk_raw["model"],
                    choices=choices
                )
            except json.JSONDecodeError:
                continue


class Chat:
    def __init__(self, client):
        self.completions = ChatCompletions(client)
