"""继续消息持久化 + 交替历史单测：POST /continue 接受的用户
消息必须落盘（与状态改回 queued 同一事务），GET
/tasks/{id}/chat-history 按时间交错拼出 prompt / 继续 /
后续消息 / 链路助手消息；409 的继续绝不能留 phantom 消息；
部署前发生的继续在 audit 里只有任务 id、正文无法追回，
接口要如实报 legacy_continue_count，不编造。

不跑真实模型回合、不联网：server 侧走本地临时端口 + 临时
SQLite，链路通过 server.trace_dir 注入临时 jsonl。
"""

import http.client
import json
import os
import sqlite3
import sys
import tempfile
import threading
import time
import unittest
from pathlib import Path

try:
    from cloud.auth import AuthStore
    from cloud.server import create_server
except ImportError:  # discover -s cloud/tests 时把项目根目录补进 sys.path
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.auth import AuthStore
    from cloud.server import create_server


class ChatHistoryTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")
        self.trace_dir = os.path.join(self.tmp.name, "traces")
        os.makedirs(self.trace_dir)
        self.auth = AuthStore(self.db_path)
        self.owner_id, self.owner_key = self.auth.create_user("owner")
        self.server = create_server(self.db_path, "127.0.0.1", 0)
        self.server.trace_dir = self.trace_dir
        self.port = self.server.server_address[1]
        self.thread = threading.Thread(
            target=self.server.serve_forever, daemon=True)
        self.thread.start()
        status, self.project = self._req(
            "POST", "/projects", {"name": "demo"}, self.owner_key)
        self.assertEqual(status, 201)

    def tearDown(self):
        self.server.shutdown()
        self.server.server_close()
        self.thread.join(timeout=5)
        try:
            self.auth.conn.close()
        except Exception:
            pass
        self.tmp.cleanup()

    def _req(self, method, path, body=None, key=None):
        conn = http.client.HTTPConnection(
            "127.0.0.1", self.port, timeout=10)
        hdrs = {}
        payload = None
        if key:
            hdrs["X-API-Key"] = key
        if body is not None:
            payload = json.dumps(body)
            hdrs["Content-Type"] = "application/json"
        conn.request(method, path, body=payload, headers=hdrs)
        resp = conn.getresponse()
        raw = resp.read()
        conn.close()
        try:
            return resp.status, json.loads(raw.decode("utf-8"))
        except (ValueError, UnicodeDecodeError):
            return resp.status, {}

    def _task(self, status="completed"):
        _s, task = self._req(
            "POST", "/tasks",
            {"project_id": self.project["id"], "prompt": "original"},
            self.owner_key)
        self.assertEqual(_s, 201, task)
        if status != "queued":
            conn = sqlite3.connect(self.db_path)
            conn.execute("UPDATE tasks SET status = ? WHERE id = ?",
                         (status, task["id"]))
            conn.commit()
            conn.close()
        return task["id"]

    def _set_status(self, task_id, status):
        conn = sqlite3.connect(self.db_path)
        conn.execute("UPDATE tasks SET status = ? WHERE id = ?",
                     (status, task_id))
        conn.commit()
        conn.close()

    def _chat(self, task_id, key=None):
        return self._req(
            "GET", "/tasks/%s/chat-history" % task_id,
            key=key or self.owner_key)

    def test_continue_persisted_and_interleaved(self):
        tid = self._task("completed")
        s, _t = self._req(
            "POST", "/tasks/%s/continue" % tid,
            {"message": "first follow-up"}, self.owner_key)
        self.assertEqual(s, 200)
        self._set_status(tid, "completed")
        s, _t = self._req(
            "POST", "/tasks/%s/continue" % tid,
            {"message": "second follow-up"}, self.owner_key)
        self.assertEqual(s, 200)
        s, chat = self._chat(tid)
        self.assertEqual(s, 200, chat)
        self.assertEqual(chat["continue_messages_stored"], 2)
        # 两次继续都有 audit，全部已入库，缺口为 0。
        self.assertEqual(chat["legacy_continue_count"], 0)
        texts = [(m["role"], m["source"], m["text"])
                 for m in chat["messages"]]
        self.assertEqual(texts, [
            ("user", "prompt", "original"),
            ("user", "continue", "first follow-up"),
            ("user", "continue", "second follow-up"),
        ])
        # 注入链路：把两次继续与任务创建时间钉到确定时间戳，
        # 助手回复落在两次继续之间与之后，按时间交错（不靠
        # 墙钟间隔赌顺序——两次继续在测试里只隔几十毫秒，
        # 助手 +500ms 本就该排在第二次继续之后）。
        conn = sqlite3.connect(self.db_path)
        conn.execute("UPDATE tasks SET created_at = ? WHERE id = ?",
                     (1699999900.0, tid))
        conn.execute(
            "UPDATE task_continue_messages SET created_at = ?"
            " WHERE task_id = ? AND body = ?",
            (1700000000.0, tid, "first follow-up"))
        conn.execute(
            "UPDATE task_continue_messages SET created_at = ?"
            " WHERE task_id = ? AND body = ?",
            (1700000100.0, tid, "second follow-up"))
        conn.commit()
        conn.close()
        entries = [
            {"kind": "message", "id": "m1", "ts": 1700000050000,
             "phase": "final_answer", "text": "done first"},
            {"kind": "message", "id": "m2", "ts": 1700000150000,
             "phase": "final_answer", "text": "done second"},
            {"kind": "exec", "id": "e1", "ts": 1700000150600,
             "command": "ls", "exitCode": 0, "durationMs": 1,
             "cwd": "/tmp"},
        ]
        with open(os.path.join(self.trace_dir, tid + ".jsonl"),
                  "w", encoding="utf-8") as f:
            for e in entries:
                f.write(json.dumps(e) + "\n")
        s, chat = self._chat(tid)
        self.assertEqual(s, 200, chat)
        texts = [(m["role"], m["source"], m["text"])
                 for m in chat["messages"]]
        self.assertEqual(texts, [
            ("user", "prompt", "original"),
            ("user", "continue", "first follow-up"),
            ("assistant", "trace", "[final_answer] done first"),
            ("user", "continue", "second follow-up"),
            ("assistant", "trace", "[final_answer] done second"),
            ("tool", "trace", "ls（exit=0，1ms，cwd=/tmp）"),
        ])

    def test_followup_and_legacy_gap_reported(self):
        # 后续消息仅支持进行中任务，completed 会 409。
        tid = self._task("running")
        s, _f = self._req(
            "POST", "/tasks/%s/followups" % tid,
            {"body": "ping while running", "mode": "direct"},
            self.owner_key)
        self.assertEqual(s, 201, _f)
        # 模拟部署前的一次继续：audit 里有记录、表里没有正文。
        conn = sqlite3.connect(self.db_path)
        conn.execute(
            "INSERT INTO audit_log(user_id, project_id, action,"
            " detail, created_at) VALUES(?, ?, 'task_continued', ?, ?)",
            (self.owner_id, self.project["id"], tid, time.time()))
        conn.commit()
        conn.close()
        s, chat = self._chat(tid)
        self.assertEqual(s, 200, chat)
        self.assertEqual(chat["continue_messages_stored"], 0)
        self.assertEqual(chat["legacy_continue_count"], 1)
        sources = [(m["role"], m["source"]) for m in chat["messages"]]
        self.assertEqual(sources, [
            ("user", "prompt"), ("user", "followup")])

    def test_rejected_continue_stores_nothing(self):
        tid = self._task("queued")
        code, _b = self._req(
            "POST", "/tasks/%s/continue" % tid,
            {"message": "should not persist"}, self.owner_key)
        self.assertEqual(code, 409)
        s, chat = self._chat(tid)
        self.assertEqual(s, 200, chat)
        self.assertEqual(chat["continue_messages_stored"], 0)
        self.assertEqual(len(chat["messages"]), 1)  # 只有 prompt

    def test_chat_history_auth(self):
        tid = self._task("completed")
        code, _b = self._chat(tid, key="bad-key")
        self.assertEqual(code, 401)
        code, _b = self._chat("no-such-task")
        self.assertEqual(code, 404)


if __name__ == "__main__":
    unittest.main()