import asyncio
import base64
from typing import Callable

from openai import AsyncOpenAI

from .config import AICfg


class AIClient:
    def __init__(self, cfg: AICfg):
        kwargs: dict = {"api_key": cfg.api_key or "sk-no-key"}
        if cfg.base_url:
            kwargs["base_url"] = cfg.base_url
        self._client = AsyncOpenAI(**kwargs)
        self._cfg = cfg

    async def analyze(self, image_path: str, on_chunk: Callable[[str], None]) -> None:
        with open(image_path, "rb") as f:
            img_b64 = base64.b64encode(f.read()).decode()
        data_url = f"data:image/png;base64,{img_b64}"

        stream = await self._client.chat.completions.create(
            model=self._cfg.model,
            max_tokens=self._cfg.max_tokens,
            stream=True,
            messages=[
                {"role": "system", "content": self._cfg.system_prompt},
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": data_url}},
                        {"type": "text", "text": "请分析截图中的题目，给出答案和解析。"},
                    ],
                },
            ],
        )

        async for chunk in stream:
            delta = chunk.choices[0].delta.content
            if delta:
                on_chunk(delta)

    def analyze_sync(self, image_path: str, on_chunk: Callable[[str], None]) -> None:
        asyncio.run(self.analyze(image_path, on_chunk))