# -*- coding: utf-8 -*-
"""
豆包（火山方舟）API 封装
使用 OpenAI 兼容接口
"""
import os
import json
import requests


class DoubaoService:
    """豆包流式对话服务"""

    def __init__(self):
        self.api_key = os.getenv("DOUBAO_API_KEY", "")
        self.model = os.getenv("DOUBAO_MODEL", "doubao-1-5-pro-32k-250115")
        self.base_url = os.getenv(
            "DOUBAO_BASE_URL", "https://ark.cn-beijing.volces.com/api/v3"
        )

    def stream_chat(self, messages):
        """
        流式对话，逐块 yield 文本内容
        :param messages: [{"role": "user"|"assistant"|"system", "content": str}]
        """
        if not self.api_key:
            raise RuntimeError("未配置 DOUBAO_API_KEY，请在 .env 中设置")

        url = f"{self.base_url}/chat/completions"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        body = {
            "model": self.model,
            "messages": messages,
            "stream": True,
        }

        with requests.post(
            url, headers=headers, json=body, stream=True, timeout=120
        ) as resp:
            if resp.status_code != 200:
                raise RuntimeError(
                    f"豆包接口错误 {resp.status_code}: {resp.text}"
                )
            for line in resp.iter_lines(decode_unicode=True):
                if not line:
                    continue
                if line.startswith("data:"):
                    data = line[5:].strip()
                    if data == "[DONE]":
                        break
                    try:
                        obj = json.loads(data)
                        delta = obj.get("choices", [{}])[0].get("delta", {})
                        content = delta.get("content")
                        if content:
                            yield content
                    except json.JSONDecodeError:
                        continue