"""任务实时链路（服务端）单测：GET /tasks/{id}/trace 与 /trace/stream。

覆盖：鉴权（401/403/404）、快照去重与坏行跳过、worktree 在别机
时的空链路 + 说明、已结束任务的流（snapshot 后立刻 done 关流）。
不跑真实模型，不启动 worker；链路目录经 server.trace_dir 注入。
"""

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

try:
    from cloud.auth import AuthStore
    from cloud.server import create_server
except ImportError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.auth import AuthStore
    from cloud.server import create_server


class TaskTraceTestCase(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, "host-traces")
        os.makedirs(self.trace_dir, exist_ok=True)
        self.auth = AuthStore(self.db_path)
        self.owner_id, self.owner_key = self.auth.create_user("owner")
        self.other_id, self.other_key = self.auth.create_user("other")
        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)
        status, self.task = self._req(
            "POST", "/tasks",
            {"project_id": self.project["id"], "prompt": "hi"},
            self.owner_key)
        self.assertEqual(status, 201)
        self.task_id = self.task["id"]

    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)
        headers = {"Content-Type": "application/json"}
        if key:
            headers["X-API-Key"] = key
        conn.request(method, path,
                     json.dumps(body) if body is not None else None, headers)
        resp = conn.getresponse()
        raw = resp.read()
        conn.close()
        try:
            data = json.loads(raw)
        except ValueError:
            data = None
        return resp.status, data

    def _set_task(self, **fields):
        conn = sqlite3.connect(self.db_path)
        try:
            for k, v in fields.items():
                conn.execute(
                    "UPDATE tasks SET %s = ? WHERE id = ?" % k,
                    (v, self.task_id))
            conn.commit()
        finally:
            conn.close()

    def _write_trace(self, lines):
        path = os.path.join(self.trace_dir, self.task_id + ".jsonl")
        with open(path, "w", encoding="utf-8") as f:
            for line in lines:
                f.write(line + "\n")

    def test_trace_auth_and_not_found(self):
        status, _ = self._req(
            "GET", "/tasks/%s/trace" % self.task_id)
        self.assertEqual(status, 401)
        status, _ = self._req(
            "GET", "/tasks/%s/trace" % self.task_id, key=self.other_key)
        self.assertEqual(status, 403)
        status, _ = self._req(
            "GET", "/tasks/nope/trace", key=self.owner_key)
        self.assertEqual(status, 404)

    def test_trace_snapshot_dedup_and_bad_lines(self):
        self._write_trace([
            json.dumps({"kind": "exec", "id": "c1", "command": "ls",
                        "exitCode": None}),
            "this is not json",
            json.dumps({"kind": "exec", "id": "c1", "command": "ls",
                        "exitCode": 0, "durationMs": 5}),
            json.dumps({"kind": "file", "id": "f1:a.py", "path": "a.py"}),
        ])
        status, data = self._req(
            "GET", "/tasks/%s/trace" % self.task_id, key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(data["status"], "queued")
        # 同 id 后到的完成态覆盖先到的开始态，保首次顺序。
        self.assertEqual([e["id"] for e in data["entries"]],
                         ["c1", "f1:a.py"])
        self.assertEqual(data["entries"][0]["exitCode"], 0)
        self.assertEqual(data["note"], "")

    def test_trace_missing_file_is_empty(self):
        status, data = self._req(
            "GET", "/tasks/%s/trace" % self.task_id, key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(data["entries"], [])

    def test_trace_remote_worktree_host_is_empty_with_note(self):
        self._write_trace([
            json.dumps({"kind": "exec", "id": "c1", "command": "ls"})])
        self._set_task(worktree_host="some-other-host")
        status, data = self._req(
            "GET", "/tasks/%s/trace" % self.task_id, key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(data["entries"], [])
        self.assertIn("some-other-host", data["note"])

    def _read_sse(self, path, key):
        conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=10)
        headers = {}
        if key:
            headers["X-API-Key"] = key
        conn.request("GET", path, headers=headers)
        resp = conn.getresponse()
        raw = resp.read()
        conn.close()
        events = []
        for block in raw.decode("utf-8").split("\n\n"):
            for line in block.split("\n"):
                if line.startswith("data: "):
                    events.append(json.loads(line[len("data: "):]))
        return resp.status, events

    def test_stream_finished_task_snapshot_then_done(self):
        self._write_trace([
            json.dumps({"kind": "exec", "id": "c1", "command": "ls",
                        "exitCode": 0})])
        self._set_task(status="completed")
        status, events = self._read_sse(
            "/tasks/%s/trace/stream" % self.task_id, self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual([e["type"] for e in events], ["snapshot", "done"])
        self.assertEqual(events[0]["entries"][0]["id"], "c1")
        self.assertEqual(events[1]["status"], "completed")

    def test_stream_auth(self):
        status, events = self._read_sse(
            "/tasks/%s/trace/stream" % self.task_id, None)
        self.assertEqual(status, 401)
        self.assertEqual(events, [])


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