# -*- coding: utf-8 -*-
"""中途切会话回归：agent 流式进行中切走 → 切回 → 时间线必须完整
（用户报告：切走再切回，正文/思考那一段显示不出来）
运行: python tests/smoke_midswitch.py
"""
import json
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
os.environ.setdefault("HAOCODE_RENDER", "software")
os.environ["QTWEBENGINE_CHROMIUM_FLAGS"] = "--disable-gpu"  # 绕过 AMD 核显 context lost

from PyQt6.QtWidgets import QApplication  # noqa: E402
from PyQt6.QtCore import QTimer  # noqa: E402
from ui.views.main_window import MainWindow  # noqa: E402

app = QApplication(sys.argv)
window = None
results = []
S = {}  # 测试状态


def check(name, ok, detail=""):
    print(f"  {'PASS' if ok else 'FAIL'}  {name}"
          + (f"  [{detail}]" if detail and not ok else ""))
    results.append(ok)


CHECK_JS = r"""
(function() {
    var mid = window.__msMid;
    var w = document.getElementById(mid);
    if (!w) return JSON.stringify({error: 'wrapper missing'});
    var tl = w.querySelector('.reply-content');
    var blocks = Array.prototype.map.call(tl.children, function(el) {
        return el.className.split(' ')[0];
    });
    var texts = Array.prototype.map.call(
        w.querySelectorAll('.md-segment'),
        function(s) { return s.textContent || ''; });
    var thinks = Array.prototype.map.call(
        w.querySelectorAll('.think-content'),
        function(s) { return s.textContent || ''; });
    return JSON.stringify({blocks: blocks, texts: texts, thinks: thinks,
        streaming: w.classList.contains('streaming')});
})()
"""


def fire(session, fn, *args):
    """模拟 worker 信号到达（UI 线程直接调 handler）"""
    getattr(window, fn)(*([session] + list(args)))


def step1():
    global window
    window = MainWindow()
    window.show()
    QTimer.singleShot(3000, step2)


def step2():
    db = window.db
    sessA = db.create_session("切走测试A")
    sessB = db.create_session("切走测试B")
    A, B = sessA["id"], sessB["id"]
    db.add_message(B, "user", "B 的问题", None)

    # ---- 会话 A 开始流式（手动模拟事件到达）----
    user_row = db.add_message(A, "user", "A 的问题", None)
    mid = "msg-midswitch"
    window.current_session_id = A
    window.chat_bridge.create_message(mid, "assistant", "", "SW")
    window._active_streams[A] = {
        "msg_id": mid, "content": "", "reasoning": "",
        "timeline": [], "tl_kind": None,
        "parent_id": user_row["id"], "branch_info": None, "worker": None,
        "previous_leaf_id": user_row["id"],
    }
    S.update(A=A, B=B, mid=mid)
    window.browser.page().runJavaScript(f"window.__msMid = '{mid}';")

    # 阶段 1（A 前台）：think1 + text1 + 工具
    fire(A, "on_reasoning_received", "思考第一")
    fire(A, "on_reasoning_received", "段内容")
    fire(A, "on_chunk_received", "工具前正文")
    fire(A, "_on_tool_started", "call-ms", "bash", '{"command":"echo ms"}')

    QTimer.singleShot(800, switch_away)


def switch_away():
    """切到 B：A 的流继续在后台跑"""
    print("[test] 切换到 B ...", flush=True)
    window.load_messages_to_web(S["B"])
    # 阶段 2（A 后台）：think2 + text2（只应累积，不进 JS）
    QTimer.singleShot(600, lambda: (
        fire(S["A"], "on_reasoning_received", "思考第二"),
        fire(S["A"], "on_chunk_received", "工具后正文"),
        print("[test] 后台 token 已累积", flush=True),
        QTimer.singleShot(600, switch_back)
    ))


def switch_back():
    """切回 A：应重建 wrapper + 恢复时间线 + 后续无缝续流"""
    print("[test] 切回 A ...", flush=True)
    window.load_messages_to_web(S["A"])
    # ★ 切回后立即检查：恢复的正文/思考必须已同步上屏（不等 rAF/timer）
    def check_right_away():
        window.browser.page().runJavaScript(CHECK_JS, lambda res: (
            print("[切回后立即]", str(res), flush=True),
            check("切回后立即: 后台正文已同步上屏",
                  "工具后正文" in json.loads(str(res)).get("texts", [""] + [json.loads(str(res))["texts"][0] if json.loads(str(res)).get("texts") else ""])[0] or "工具后正文" in "".join(json.loads(str(res)).get("texts", []))),
            QTimer.singleShot(300, phase3)
        ))
    def phase3():
        fire(S["A"], "on_chunk_received", " 续流正文")
        QTimer.singleShot(600, finish_it)
    QTimer.singleShot(400, check_right_away)


def finish_it():
    window.on_reply_finished(S["A"])
    QTimer.singleShot(400, step_check)


DIAG_JS = r"""
(function() {
    var wrappers = Array.prototype.map.call(
        document.querySelectorAll('.message-wrapper'),
        function(w) { return w.id + ':' + w.className.split(' ')[1]; });
    return JSON.stringify({wrappers: wrappers,
        bufKeys: Object.keys(window.messageBuffer || {})});
})()
"""

def step_check():
    def diag(res):
        print("[diag]", str(res), flush=True)
        window.browser.page().runJavaScript(CHECK_JS, on_check)
    window.browser.page().runJavaScript(DIAG_JS, diag)


def on_check(res):
    d = json.loads(str(res))
    if "error" in d:
        check("切回后 wrapper 存在", False, d["error"])
        done()
        return
    blocks = d["blocks"]
    texts = "".join(d["texts"])
    thinks = "".join(d["thinks"])
    print(f"  blocks = {blocks}")
    print(f"  texts  = {d['texts']}")
    print(f"  thinks = {d['thinks']}")
    check("切回后 wrapper 存在", True)
    check("块顺序与事件顺序一致: think→text→chip→think→text",
          [b for b in blocks if b in
           ("think-block", "tool-chip", "md-segment")] ==
          ["think-block", "md-segment", "tool-chip", "think-block", "md-segment"],
          str(blocks))
    check("正文完整（切走前+后台+切回后）",
          "工具前正文" in texts and "工具后正文" in texts
          and "续流正文" in texts, str(d["texts"]))
    check("思考完整（切走前+后台）",
          "思考第一段内容" in thinks and "思考第二" in thinks,
          str(d["thinks"]))
    check("streaming 已收尾", not d["streaming"])
    # 入库验证
    row = window.db.get_message_chain(S["A"])
    asst = [m for m in row if m["role"] == "assistant"]
    ok_db = False
    if asst:
        tl = asst[-1].get("timeline")
        if tl:
            tl_list = json.loads(tl)
            joined = "".join(e.get("text", "") for e in tl_list if e["t"] == "text")
            ok_db = ("工具前正文" in joined and "工具后正文" in joined
                     and "续流正文" in joined)
    check("DB 时间线完整（三段正文都在）", ok_db)
    done()


def done():
    try:
        window.db.delete_session(S["A"])
        window.db.delete_session(S["B"])
    except Exception:
        pass
    print(f"\n===== {'ALL PASS' if all(results) else 'HAS FAILURES'}: "
          f"{sum(results)}/{len(results)} =====")
    app.quit()


QTimer.singleShot(500, step1)
app.exec()