import requests
import json


CHAT_COMPLETIONS_URL = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
BOTS_CHAT_COMPLETIONS_URL = "https://ark.cn-beijing.volces.com/api/v3/bots/chat/completions"
API_KEY = "Bearer fcc424e5-58af-494d-9683-5787413a26c9"
DEFAULT_BATCH_SIZE = 50


def _streaming_request(url, model,messages,type, batch_size=50):
    payload = {
        "model": model,
        "stream": True,
        "messages": messages,
        "thinking":{
            "type":type
        }
    }

    headers = {
        'Authorization': API_KEY,
        'Content-Type': 'application/json'
    }
    try:
        response = requests.post(
            url,
            json=payload,
            headers=headers,
            stream=True,
            timeout=300
        )
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        err_msg = f"请求失败: {str(e)}"

        def error_generator():
            error_data = json.dumps({"message": err_msg, "type": 0}, ensure_ascii=False)
            yield f"data: {error_data}\n\n"

        return error_generator()





    def generate_sse():
        buffer = ""

        for line in response.iter_lines():
            if not line:
                continue
            try:
                line_str = line.decode('utf-8').lstrip('data: ').strip()
                if line_str == '[DONE]':
                    break

                chunk = json.loads(line_str)

                if isinstance(chunk, dict) and "choices" in chunk:
                    choice = chunk["choices"][0] if chunk["choices"] else None
                    if choice and "delta" in choice:
                        delta = choice["delta"]

                        reasoning_content = delta.get("reasoning_content")
                        if reasoning_content:

                            reasoning_data = json.dumps({
                                "message": reasoning_content,
                                "type": 4
                            }, ensure_ascii=False)
                            yield f"data: {reasoning_data}\n\n"


                        content = delta.get("content")
                        if content:
                            buffer += content
                            if len(buffer) >= batch_size:
                                data = json.dumps({
                                    "message": buffer,
                                    "type": 1
                                }, ensure_ascii=False)
                                yield f"data: {data}\n\n"
                                buffer = ""

            except json.JSONDecodeError:
                continue
            except Exception as e:
                error_data = json.dumps({
                    "message": f"处理数据错误: {str(e)}",
                    "type": 0
                }, ensure_ascii=False)
                yield f"data: {error_data}\n\n"
                break
        if buffer:
            data = json.dumps({"message": buffer, "type": 1},
                              ensure_ascii=False)
            yield f"data: {data}\n\n"

    return generate_sse()


def _streaming_request_comment(url, model,messages,type, batch_size=50):
    payload = {
        "model": model,
        "stream": True,
        "messages": messages,
        "thinking":{
            "type":type
        }
    }
    headers = {
        'Authorization': API_KEY,
        'Content-Type': 'application/json'
    }
    try:
        response = requests.post(
            url,
            json=payload,
            headers=headers,
            stream=True,
            timeout=300
        )
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        err_msg = f"请求失败: {str(e)}"

        def error_generator():
            error_data = json.dumps({"message": err_msg, "type": 0}, ensure_ascii=False)
            yield f"data: {error_data}\n\n"

        return error_generator()


    def generate_sse():
        buffer = ""

        for line in response.iter_lines():
            if not line:
                continue
            try:
                line_str = line.decode('utf-8').lstrip('data: ').strip()
                if line_str == '[DONE]':
                    break

                chunk = json.loads(line_str)

                if isinstance(chunk, dict) and "choices" in chunk:
                    choice = chunk["choices"][0] if chunk["choices"] else None
                    if choice and "delta" in choice:
                        delta = choice["delta"]

                        reasoning_content = delta.get("reasoning_content")
                        if reasoning_content:

                            reasoning_data = json.dumps({
                                "message": reasoning_content,
                                "type": 4
                            }, ensure_ascii=False)
                            yield f"data: {reasoning_data}\n\n"


                        content = delta.get("content")
                        if content:
                            buffer += content
                            if len(buffer) >= batch_size:
                                data = json.dumps({
                                    "message": buffer,
                                    "type": 6
                                }, ensure_ascii=False)
                                yield f"data: {data}\n\n"
                                buffer = ""

            except json.JSONDecodeError:
                continue
            except Exception as e:
                error_data = json.dumps({
                    "message": f"处理数据错误: {str(e)}",
                    "type": 0
                }, ensure_ascii=False)
                yield f"data: {error_data}\n\n"
                break
        if buffer:
            data = json.dumps({"message": buffer, "type": 6},
                              ensure_ascii=False)
            yield f"data: {data}\n\n"

    return generate_sse()


# 标准聊天
def get_ai_chat_history_stream(messages, type = "disabled", batch_size=DEFAULT_BATCH_SIZE):

    return _streaming_request(
        url=CHAT_COMPLETIONS_URL,
        model="ep-20250804181425-5blp4",
        messages=messages,
        type=type,
        batch_size=batch_size
    )


# 联网聊天
def get_bot_with_history_stream(messages,type = "auto", batch_size=DEFAULT_BATCH_SIZE):
    return _streaming_request(
        url=BOTS_CHAT_COMPLETIONS_URL,
        model="bot-20250730164611-nbsng",
        messages=messages,
        type=type,
        batch_size=batch_size
    )


# 知识库
def get_bot_with_history_rag_stream(messages,type = "auto", batch_size=DEFAULT_BATCH_SIZE):
    return _streaming_request(
        url=BOTS_CHAT_COMPLETIONS_URL,
        model="bot-20250725180737-mjxbx",
        messages=messages,
        type=type,
        batch_size=batch_size
    )


# 图像识别
def get_image_to_text_stream(content,type = "disabled", batch_size=DEFAULT_BATCH_SIZE):
    messages = [{"content": content, "role": "user"}]
    return _streaming_request(
        url=CHAT_COMPLETIONS_URL,
        model="doubao-seed-2-0-lite-260215",
        messages=messages,
        type=type,
        batch_size=batch_size
    )


def get_ai_chat_history(messages, stream=False):
    url = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"

    payload = json.dumps({
        "model": "ep-20241214232538-x27zp",
        "stream": stream,
        "messages": messages
    })
    headers = {
        'Authorization': API_KEY,
        'Content-Type': 'application/json'
    }
    response = requests.request("POST", url, headers=headers, data=payload)
    data = response.json()
    ai_reply = data.get('choices', [{}])[0].get('message', {})
    total_tokens = data.get('usage', {}).get('total_tokens')
    return ai_reply, total_tokens


def get_ai_chat_history_stream_comment(messages, type = "disabled", batch_size=DEFAULT_BATCH_SIZE):

    return _streaming_request_comment(
        url=CHAT_COMPLETIONS_URL,
        model="ep-20250804181425-5blp4",
        messages=messages,
        type=type,
        batch_size=batch_size
    )



