import base64
import mimetypes
from dataclasses import dataclass
from pathlib import Path
from typing import Literal, cast
from pydantic import BaseModel

import httpx2
from openai import OpenAI
from openai.types.chat import (
    ChatCompletionContentPartParam,
    ChatCompletionMessageParam,
)


@dataclass
class AgentMessage:
    content: str
    """Content of the message"""

    role: Literal["system", "assistant", "user"]
    """Who sent the message"""

class Agent:
    """Perform operations with timeline events using OpenAI-compatible API"""
    def __init__(self, *, model: str, base_url: str | None, api_key: str, **kwargs) -> None:
        self._client = OpenAI(
            base_url=base_url,
            api_key=api_key,
            http_client=httpx2.Client(verify=False),
            **kwargs
        )
        self._model = model
        

    def completion(self, messages: list[AgentMessage], **kwargs) -> str:
        """Generate a completion for specified messages."""
        messages_raw: list[ChatCompletionMessageParam] = []
        for m in messages:
            messages_raw.append(
                cast(ChatCompletionMessageParam, {
                    "role": m.role,
                    "content": m.content
                })
            )
        response = self._client.chat.completions.parse(
            model=self._model,
            messages=messages_raw,
            **kwargs
        )
        return response.choices[0].message.content # type: ignore

    def completion_with_images(
        self,
        messages: list[AgentMessage],
        image_paths: list[str],
        *,
        detail: Literal["low", "high", "auto"] = "auto",
        **kwargs,
    ) -> str:
        """Generate a completion with local images attached to the last message."""
        if not messages or messages[-1].role != "user":
            raise ValueError("The last message must be a user message")
        if not image_paths:
            raise ValueError("At least one image is required")

        messages_raw: list[ChatCompletionMessageParam] = [
            cast(ChatCompletionMessageParam, {
                "role": message.role,
                "content": message.content,
            })
            for message in messages[:-1]
        ]
        content: list[ChatCompletionContentPartParam] = [
            {"type": "text", "text": messages[-1].content}
        ]
        for index, image_path in enumerate(image_paths, start=1):
            path = Path(image_path)
            mime_type = mimetypes.guess_type(path.name)[0] or "image/jpeg"
            if not mime_type.startswith("image/"):
                raise ValueError(f"Unsupported image type: {image_path}")
            encoded = base64.b64encode(path.read_bytes()).decode("ascii")
            content.append({
                "type": "text",
                "text": f"Изображение {index}",
            })
            content.append({
                "type": "image_url",
                "image_url": {
                    "url": f"data:{mime_type};base64,{encoded}",
                    "detail": detail,
                },
            })
        messages_raw.append({
            "role": "user",
            "content": content,
        })

        response = self._client.chat.completions.parse(
            model=self._model,
            messages=messages_raw,
            **kwargs,
        )
        return response.choices[0].message.content # type: ignore