from __future__ import annotations

import json
from abc import ABC, abstractmethod
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Optional, TypedDict
from urllib import error as urllib_error
from urllib import request as urllib_request

from ramalama.chat_utils import ChatMessageType
from ramalama.config import ActiveConfig


class RequiredChatRequestOptionsDict(TypedDict):
    stream: bool


class ChatRequestOptionsDict(RequiredChatRequestOptionsDict, total=False):
    model: str
    temperature: float
    max_tokens: int


@dataclass
class ChatRequestOptions:
    """Normalized knobs for building a chat completion request."""

    model: Optional[str]
    stream: bool = True
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None
    extra: Optional[dict[str, Any]] = None

    def to_dict(self) -> ChatRequestOptionsDict:
        result: ChatRequestOptionsDict = {'stream': self.stream}

        if self.temperature is not None:
            result['temperature'] = self.temperature
        if self.max_tokens is not None:
            result['max_tokens'] = self.max_tokens
        if self.model is not None:
            result['model'] = self.model

        if self.extra is not None:
            # Can't easily type this without access to open TypedDicts
            # https://typing.python.org/en/latest/spec/glossary.html#term-closed
            result |= self.extra  # type: ignore

        return result


@dataclass
class ChatStreamEvent:
    """A provider-agnostic representation of a streamed delta."""

    text: Optional[str] = None
    raw: Optional[dict[str, Any]] = None
    done: bool = False


class ChatProviderError(Exception):
    """Raised when a provider request fails or returns an invalid payload."""

    def __init__(self, message: str, *, status_code: Optional[int] = None, payload: Optional[Any] = None):
        super().__init__(message)
        self.status_code = status_code
        self.payload = payload


class ChatProvider(ABC):
    """Abstract base class for hosted chat providers."""

    provider: str = "base"
    default_path: str

    def __init__(
        self,
        base_url: str,
        api_key: Optional[str] = None,
        default_headers: Optional[Mapping[str, str]] = None,
    ) -> None:
        if api_key is None:
            api_key = ActiveConfig().api_key

        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self._default_headers: dict[str, str] = dict(default_headers or {})

    def build_url(self, path: Optional[str] = None) -> str:
        rel = path or self.default_path
        if not rel.startswith("/"):
            rel = f"/{rel}"
        return f"{self.base_url}{rel}"

    def prepare_headers(
        self,
        *,
        include_auth: bool = True,
        extra: Optional[dict[str, str]] = None,
        options: Optional[ChatRequestOptions] = None,
    ) -> dict[str, str]:
        headers: dict[str, str] = {
            "Content-Type": "application/json",
            **self._default_headers,
            **self.provider_headers(options),
        }

        if include_auth:
            headers.update(self.auth_headers())
        if extra:
            headers.update(extra)
        return headers

    def auth_headers(self) -> dict[str, str]:
        return {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}

    def serialize_payload(self, payload: Mapping[str, Any]) -> bytes:
        return json.dumps(payload).encode("utf-8")

    def create_request(
        self, messages: Sequence[ChatMessageType], options: ChatRequestOptions
    ) -> urllib_request.Request:
        payload = self.build_payload(messages, options)
        headers = self.prepare_headers(options=options, extra=self.additional_request_headers(options))
        body = self.serialize_payload(payload)
        return urllib_request.Request(
            self.build_url(self.resolve_request_path(options)),
            data=body,
            headers=headers,
            method="POST",
        )

    # ------------------------------------------------------------------
    # Provider customization points
    # ------------------------------------------------------------------
    def provider_headers(self, options: Optional[ChatRequestOptions] = None) -> dict[str, str]:
        return {}

    def additional_request_headers(self, options: Optional[ChatRequestOptions] = None) -> dict[str, str]:
        return {}

    def resolve_request_path(self, options: Optional[ChatRequestOptions] = None) -> str:
        return self.default_path

    @abstractmethod
    def build_payload(self, messages: Sequence[ChatMessageType], options: ChatRequestOptions) -> Mapping[str, Any]:
        """Return the provider-specific payload."""

    @abstractmethod
    def parse_stream_chunk(self, chunk: bytes) -> Iterable[ChatStreamEvent]:
        """Yield zero or more events parsed from a streamed response chunk."""

    # ------------------------------------------------------------------
    # Error handling
    # ------------------------------------------------------------------
    def raise_for_status(self, status_code: int, payload: Optional[Any] = None) -> None:
        if status_code >= 400:
            if isinstance(payload, dict) and "error" in payload:
                err = payload["error"]
                message = str(err.get("message") or err.get("type") or err) if isinstance(err, dict) else str(err)
            else:
                message = "chat request failed"

            raise ChatProviderError(message, status_code=status_code, payload=payload)

    # ------------------------------------------------------------------
    # Non-streamed helpers
    # ------------------------------------------------------------------
    def parse_response_body(self, body: bytes) -> Any:
        if not body:
            return None
        return json.loads(body.decode("utf-8"))

    def list_models(self) -> list[str]:
        """Return available model identifiers exposed by the provider."""

        request = urllib_request.Request(
            self.build_url("/models"),
            headers=self.prepare_headers(include_auth=True),
            method="GET",
        )
        try:
            with urllib_request.urlopen(request) as response:
                payload = self.parse_response_body(response.read())
        except urllib_error.HTTPError as exc:
            if exc.code in (401, 403):
                message = (
                    f"Could not authenticate with {self.provider}."
                    " The provided API key was either missing or invalid.\n"
                    f"Set RAMALAMA_API_KEY or ramalama.provider.<provider_name>.api_key."
                )
                try:
                    payload = self.parse_response_body(exc.read())
                except Exception:
                    payload = {}

                if details := payload.get("error", {}).get("message", None):
                    message = f"{message}\n\n{details}"

                raise ChatProviderError(message, status_code=exc.code) from exc
            raise

        if not isinstance(payload, Mapping):
            raise ChatProviderError("Invalid model list payload", payload=payload)

        data = payload.get("data")
        if not isinstance(data, list):
            raise ChatProviderError("Invalid model list payload", payload=payload)

        models: list[str] = []
        for entry in data:
            if isinstance(entry, Mapping) and (model_id := entry.get("id")):
                models.append(str(model_id))

        return models


__all__ = [
    "ChatProvider",
    "ChatProviderError",
    "ChatRequestOptions",
    "ChatStreamEvent",
]
