entry.py

"""entry.py：simplified_chat_engine 的 OpenAI 兼容中转站（Flask）。

QQPilot（C#）把 server_url 指向本服务（如 http://localhost:7749/v1）后，
群聊请求会进入这里，由 answer.getAnswer 跑完整流水线
（EcoReasoning → clumsyImitation → forgetfulMem → Xreply）。

超时策略：EcoReasoning 判定"不需要回复"时，getAnswer 返回 None，
本服务会 sleep config.timeout 秒（默认 300，对应 QQPilot 的 remote_server_timeout），
让 QQPilot 端请求超时、放弃回复——即"不需要回复就等到 QQPilot 超时"。
"""

import hashlib
import os
import time
from typing import List, Optional

from flask import Flask, jsonify, request

import config as engine_config
from answer import GetAnswer
from chatContent import ChatContent

app = Flask(__name__)

# http(s) 图片下载缓存目录（本地模型访问不了远程 URL 时先落盘）
_cache_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".engine_cache", "images")
os.makedirs(_cache_dir, exist_ok=True)


def _download_image(url: str) -> Optional[str]:
    """下载 http(s) 图片到本地缓存（本地模型访问不了远程 URL 时用）。失败返回 None。"""
    try:
        import httpx
        name = hashlib.sha1(url.encode("utf-8")).hexdigest()[:20] + ".jpg"
        path = os.path.join(_cache_dir, name)
        if not os.path.exists(path):
            r = httpx.get(url, timeout=10, follow_redirects=True)
            if r.status_code == 200:
                with open(path, "wb") as f:
                    f.write(r.content)
            else:
                return None
        return path
    except Exception as e:
        print(f"× 图片下载失败: {e}")
        return None


def _parse_content(content) -> tuple[str, List[str], List[str]]:
    """把 OpenAI 消息的 content（str 或分段数组）解析为 (text, imagePaths, imageDataUrls)。

    data URL（QQPilot 只发这种）直接保留为 imageDataUrls，不落盘；
    http(s) URL 本地模型访问不了，先下载到缓存目录再作为 imagePaths。
    """
    if isinstance(content, str):
        return content, [], []
    if isinstance(content, list):
        text_parts: List[str] = []
        img_paths: List[str] = []
        img_data_urls: List[str] = []
        for part in content:
            if not isinstance(part, dict):
                continue
            ptype = part.get("type")
            if ptype == "text":
                text_parts.append(part.get("text", "") or "")
            elif ptype == "image_url":
                url = (part.get("image_url") or {}).get("url", "") if isinstance(part.get("image_url"), dict) else ""
                if url.startswith("data:"):
                    img_data_urls.append(url)
                elif url.startswith("http://") or url.startswith("https://"):
                    path = _download_image(url)
                    if path:
                        img_paths.append(path)
        return "".join(text_parts), img_paths, img_data_urls
    return "", [], []


def _messages_to_chat_contents(messages: list) -> tuple[Optional[str], List[ChatContent]]:
    """OpenAI messages → (system, List[ChatContent])。第一条 system 作为 systemPrompt。"""
    system: Optional[str] = None
    contents: List[ChatContent] = []
    for message in messages:
        role = message.get("role")
        content = message.get("content")
        if role == "system" and system is None:
            system = content if isinstance(content, str) else ""
            continue
        text, img_paths, img_data_urls = _parse_content(content)
        # 发消息的人：user 用 username 占位，assistant 视为机器人自己
        username = "User" if role == "user" else "Assistant"
        own = role == "assistant"
        # 保留 role 信息，避免同名歧义：assistant → ownByMyself=True
        contents.append(ChatContent(
            username=username, imagePaths=img_paths, text=text,
            time=time.strftime("%m-%d %H:%M:%S"), ownByMyself=own,
            imageDataUrls=img_data_urls))
    return system, contents


def _make_response(model: str, content: str) -> dict:
    return {
        "id": f"chatcmpl-engine-{int(time.time())}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": content},
            "finish_reason": "stop",
        }],
        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
    }


@app.route('/v1/chat/completions', methods=['POST'])
def chat_completions():
    body = request.get_json(force=True, silent=True)
    if not body or 'model' not in body:
        return jsonify({"error": "Missing 'model' in request body"}), 400
    if 'messages' not in body or not body['messages']:
        return jsonify({"error": "Missing 'messages' in request body"}), 400

    cfg = engine_config.get_config()
    model = body["model"]
    system, contents = _messages_to_chat_contents(body["messages"])

    print("\n" + "=" * 70)
    print(f"🟢 收到请求 model={model} messages={len(contents)}")
    print("=" * 70)

    t0 = time.time()
    reply = GetAnswer(contents, systemPrompt=system if system is not None else "auto")

    if reply is None:
        # EcoReasoning 判定不需要回复：拖到 QQPilot 超时。
        # 分块 sleep：客户端（QQPilot）断开后不会继续白等整个超时；
        # 建议 config.xml 的 <timeout> 略小于 QQPilot config.ini 的 remote_server_timeout，
        # 这样 QQPilot 先超时放弃，服务端也能尽快结束等待。
        print(f"⏳ 判定无需回复，等待 {cfg.timeout}s 让 QQPilot 超时...")
        deadline = time.time() + cfg.timeout
        while time.time() < deadline:
            time.sleep(min(5, deadline - time.time()))
        print("✅ 已等待完成，返回空回复")
        return jsonify(_make_response(model, ""))

    print(f"✅ 回复（用时 {time.time()-t0:.1f}s）: {reply!r}")
    return jsonify(_make_response(model, reply))


@app.route('/', methods=['GET'])
def home():
    return jsonify({
        "message": "SimplifiedChatEngine",
        "endpoint": "/v1/chat/completions",
        "note": "QQPilot 配置 server_url=http://localhost:7749/v1 即可对接",
    })


if __name__ == '__main__':
    cfg = engine_config.get_config()
    print("🚀 启动 SimplifiedChatEngine 中转站...")
    print("监听地址: http://localhost:7749/v1/chat/completions")
    print(f"无需回复时的等待超时: {cfg.timeout}s")
    print("请在 QQPilot 的 config.ini 中设置:")
    print("  server_url = http://localhost:7749/v1")
    print("  API_KEY = 任意值")
    print("\n等待请求中...（按 Ctrl+C 停止）\n")
    # threaded=True：每个请求独立线程，sleep 等待超时不会阻塞其他请求
    app.run(host='0.0.0.0', port=7749, debug=True, threaded=True)