import asyncio
import json
import logging
from collections.abc import AsyncGenerator

from zhipuai import ZhipuAI

from app.engine.providers.base import AIProvider

logger = logging.getLogger(__name__)

DEFAULT_MODEL = "glm-4-flash"


class GLMProvider(AIProvider):
    def __init__(self, api_key: str, model: str = "glm-4-flash", base_url: str = ""):
        self.client = ZhipuAI(api_key=api_key, base_url=base_url or None)
        self.model = model

    async def chat(self, messages: list[dict], system_prompt: str) -> str:
        full_messages = [{"role": "system", "content": system_prompt}] + messages
        return await self._call(full_messages, temperature=0.3)

    async def stream_chat(
        self, messages: list[dict], system_prompt: str
    ) -> AsyncGenerator[str, None]:
        full_messages = [{"role": "system", "content": system_prompt}] + messages

        def _sync_stream():
            resp = self.client.chat.completions.create(
                model=self.model,
                messages=full_messages,
                stream=True,
                temperature=0.3,
            )
            for chunk in resp:
                delta = chunk.choices[0].delta
                if delta and delta.content:
                    yield delta.content

        try:
            loop = asyncio.get_event_loop()
            gen = _sync_stream()

            def _next_token(g):
                try:
                    return next(g)
                except StopIteration:
                    return None

            while True:
                token = await loop.run_in_executor(None, _next_token, gen)
                if token is None:
                    break
                yield token
        except Exception:
            logger.exception("GLM streaming call failed")
            raise

    async def chat_with_extraction(
        self, messages: list[dict], system_prompt: str
    ) -> tuple[str, dict]:
        full_messages = [{"role": "system", "content": system_prompt}] + messages
        if not messages:
            full_messages.append({"role": "user", "content": "请根据以上要求输出结果。"})

        raw = await self._call(full_messages, temperature=0.1)
        extracted = self._parse_json(raw)
        return raw, extracted

    async def _call(self, messages: list[dict], temperature: float = 0.3) -> str:
        def _sync_call():
            resp = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=temperature,
            )
            return resp.choices[0].message.content

        try:
            return await asyncio.to_thread(_sync_call)
        except Exception:
            logger.exception("GLM API call failed")
            raise

    @staticmethod
    def _parse_json(text: str) -> dict:
        try:
            start = text.index("{")
            end = text.rindex("}") + 1
            return json.loads(text[start:end])
        except (ValueError, json.JSONDecodeError):
            logger.warning("Failed to parse JSON from model output")
            return {}