import json
from dataclasses import dataclass
from urllib import request
from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Union

from approach.config import ProviderConfig


PostFn = Callable[..., Any]


@dataclass(frozen=True)
class EncodedImage:
    data: str
    media_type: str = "image/jpeg"


class UrlLibResponse:
    def __init__(self, status: int, body: bytes):
        self.status_code = status
        self._body = body

    def raise_for_status(self):
        if self.status_code >= 400:
            raise RuntimeError(f"HTTP request failed with status {self.status_code}")

    def json(self):
        return json.loads(self._body.decode("utf-8"))


def default_post(url: str, headers: Mapping[str, str], json: Mapping[str, Any], timeout: int):
    req = request.Request(
        url,
        data=json_dumps_bytes(json),
        headers=dict(headers),
        method="POST",
    )
    with request.urlopen(req, timeout=timeout) as response:
        return UrlLibResponse(response.status, response.read())


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


def make_multimodal_content(
    text: str,
    image_b64s: Optional[Iterable[Union[str, EncodedImage]]] = None,
) -> List[Dict[str, Any]]:
    content = [{"type": "text", "text": text}]
    for image in image_b64s or []:
        encoded = image if isinstance(image, EncodedImage) else EncodedImage(data=image)
        content.append(
            {
                "type": "image_url",
                "image_url": {
                    "url": f"data:{encoded.media_type};base64,{encoded.data}"
                },
            }
        )
    return content


class OpenAICompatibleChatClient:
    def __init__(self, profile: ProviderConfig, post: Optional[PostFn] = None):
        self.profile = profile
        self.post = post or default_post

    def build_payload(
        self,
        prompt: str,
        image_b64s: Optional[Iterable[Union[str, EncodedImage]]] = None,
        response_format: Optional[Mapping[str, Any]] = None,
        temperature: float = 0,
        max_tokens: int = 4096,
    ) -> Dict[str, Any]:
        return self.build_messages_payload(
            [
                {
                    "role": "user",
                    "content": make_multimodal_content(prompt, image_b64s),
                }
            ],
            response_format=response_format,
            temperature=temperature,
            max_tokens=max_tokens,
        )

    def build_messages_payload(
        self,
        messages: Iterable[Mapping[str, Any]],
        response_format: Optional[Mapping[str, Any]] = None,
        temperature: float = 0,
        max_tokens: int = 4096,
    ) -> Dict[str, Any]:
        payload: Dict[str, Any] = {
            "model": self.profile.model,
            "temperature": temperature,
            "messages": [dict(message) for message in messages],
            "max_tokens": max_tokens,
        }
        if response_format:
            payload["response_format"] = dict(response_format)
        if self.profile.router_options:
            payload["provider"] = dict(self.profile.router_options)
        return payload

    def chat_completion(self, payload: Mapping[str, Any]) -> Mapping[str, Any]:
        if not self.profile.api_key:
            raise RuntimeError(f"Missing API key in environment variable {self.profile.api_key_env}")
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.profile.api_key}",
        }
        response = self.post(
            f"{self.profile.base_url.rstrip('/')}/chat/completions",
            headers=headers,
            json=dict(payload),
            timeout=120,
        )
        response.raise_for_status()
        return response.json()

    def complete_json(
        self,
        prompt: str,
        image_b64s: Optional[Iterable[Union[str, EncodedImage]]] = None,
        response_format: Optional[Mapping[str, Any]] = None,
    ) -> Any:
        payload = self.build_payload(prompt, image_b64s, response_format=response_format)
        return self.complete_json_payload(payload)

    def complete_json_messages(
        self,
        messages: Iterable[Mapping[str, Any]],
        response_format: Optional[Mapping[str, Any]] = None,
    ) -> Any:
        payload = self.build_messages_payload(
            messages,
            response_format=response_format,
        )
        return self.complete_json_payload(payload)

    def complete_json_payload(self, payload: Mapping[str, Any]) -> Any:
        response = self.chat_completion(payload)
        content = response["choices"][0]["message"]["content"].strip()
        if content.startswith("```"):
            content = content.strip("`").split("\n", 1)[-1]
        return json.loads(content)