import json, uuid, threading
from flask import Flask, request, jsonify
from camoufox.sync_api import Camoufox

app = Flask(__name__)
lock = threading.Lock()

with Camoufox(headless=True) as browser:
    ctx = browser.new_context()
    page = ctx.new_page()
    page.goto("https://chat.z.ai")
    page.wait_for_load_state("networkidle")
    print("Camoufox ready", flush=True)

    def do_chat(messages, model):
        body = json.dumps({
            "stream": True,
            "model": model,
            "messages": messages,
            "params": {},
            "extra": {},
            "features": {"image_generation": False, "web_search": False, "preview_mode": True, "flags": [], "enable_thinking": False},
            "chat_id": str(uuid.uuid4()),
            "id": str(uuid.uuid4()),
        })
        js = """(bodyStr) => {
            return new Promise((resolve) => {
                fetch("/api/v2/chat/completions", {
                    method: "POST",
                    headers: {"Content-Type": "application/json"},
                    body: bodyStr
                }).then(r => r.text()).then(t => resolve(t)).catch(e => resolve("ERROR:" + e));
                setTimeout(() => resolve("TIMEOUT"), 30000);
            });
        }"""
        result = page.evaluate(js, body)
        content = ""
        for line in str(result).splitlines():
            if line.startswith("data: "):
                try:
                    d = json.loads(line[6:])
                    delta = d.get("data", {}).get("delta_content", "")
                    phase = d.get("data", {}).get("phase", "")
                    if phase == "answer":
                        content += delta
                except:
                    pass
        return content

    @app.route("/camoufox/chat", methods=["POST"])
    def completions():
        data = request.json
        messages = data.get("messages", [])
        model = data.get("model", "glm-4.7")
        with lock:
            content = do_chat(messages, model)
        return jsonify({"content": content})

    @app.route("/camoufox/health", methods=["GET"])
    def health():
        return jsonify({"status": "ok"})

    app.run(port=8002, threaded=False)