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

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

import http.client
import json
import os
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, _tenant_trace_dir
except ImportError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.auth import AuthStore
    from cloud.server import create_server, _tenant_trace_dir


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.tenants_root = os.path.join(self.tmp.name, "tenants")
        os.makedirs(self.tenants_root, 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.server.tenants_root = self.tenants_root
        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 _write_tenant_trace(self, lines, filename=None):
        d = os.path.join(self.tenants_root, self.project["id"],
                         "work", "host-traces")
        os.makedirs(d, exist_ok=True)
        path = os.path.join(d, filename or (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, [])

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

    def test_agents_chain_groups_calls_by_agent(self):
        self._set_task(status="running")
        self._write_trace([
            json.dumps({"kind": "exec", "id": "c1", "thread": "th-main",
                        "command": "make", "ts": 1000}),
            json.dumps({"kind": "tool", "id": "t1", "thread": "th-sub",
                        "server": "s", "tool": "read", "ts": 1001}),
            json.dumps({"kind": "agent-start", "id": "agent:rev",
                        "name": "rev", "prompt": "review code",
                        "trace": "rev.jsonl", "ts": 1002}),
            json.dumps({"kind": "agent-done", "id": "agent:rev:done",
                        "name": "rev", "status": "completed",
                        "trace": "rev.jsonl", "ts": 1004}),
        ])
        with open(os.path.join(self.trace_dir, "rev.jsonl"),
                  "w", encoding="utf-8") as f:
            f.write(json.dumps({"kind": "exec", "id": "c9",
                                "command": "grep", "ts": 1003}) + "\n")
        status, data = self._req(
            "GET", "/tasks/%s/agents" % self.task_id, key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(data["status"], "running")
        by_kind = {a["kind"]: a for a in data["agents"]}
        main = by_kind["main"]
        self.assertEqual(main["status"], "running")
        self.assertEqual([e["id"] for e in main["calls"]], ["c1"])
        model = by_kind["model"]
        self.assertEqual(model["status"], "running")
        self.assertEqual([e["id"] for e in model["calls"]], ["t1"])
        delegate = [a for a in data["agents"]
                    if a["kind"] == "delegate"][0]
        self.assertEqual(delegate["name"], "rev")
        self.assertEqual(delegate["status"], "completed")
        self.assertEqual(delegate["prompt"], "review code")
        self.assertEqual([e["id"] for e in delegate["calls"]], ["c9"])
        self.assertEqual(delegate["updated_ts"], 1004)

    def test_agents_delegate_without_done_stays_running(self):
        # 委派进程与父任务相互独立：父任务已结束、无结束标记时，
        # 子任务仍算进行中（不能谎报它已中断）。
        self._set_task(status="completed")
        self._write_trace([
            json.dumps({"kind": "agent-start", "id": "agent:rev",
                        "name": "rev", "prompt": "review",
                        "trace": "rev.jsonl", "ts": 1002}),
        ])
        status, data = self._req(
            "GET", "/tasks/%s/agents" % self.task_id, key=self.owner_key)
        self.assertEqual(status, 200)
        delegate = [a for a in data["agents"]
                    if a["kind"] == "delegate"][0]
        self.assertEqual(delegate["status"], "running")
        self.assertEqual(data["agents"][0]["status"], "completed")

    def test_agents_missing_file_is_main_only(self):
        status, data = self._req(
            "GET", "/tasks/%s/agents" % self.task_id, key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(len(data["agents"]), 1)
        self.assertEqual(data["agents"][0]["kind"], "main")
        self.assertEqual(data["agents"][0]["calls"], [])

    def test_tenant_trace_dir_mapping(self):
        class S:
            tenants_root = "/tenants"
        self.assertEqual(
            _tenant_trace_dir(S(), "abc-123"),
            "/tenants/abc-123/work/host-traces")
        self.assertIsNone(_tenant_trace_dir(S(), None))
        self.assertIsNone(_tenant_trace_dir(S(), ""))
        self.assertIsNone(_tenant_trace_dir(S(), ".."))
        self.assertIsNone(_tenant_trace_dir(S(), "..."))

    def test_trace_reads_tenant_dir_only(self):
        # 租户 Docker 任务的链路文件只在租户目录里：旧版只读
        # 共享目录会误报为空，修后 /trace 与 /agents 都看得见。
        self._write_tenant_trace([
            json.dumps({"kind": "exec", "id": "c1", "command": "make"}),
        ])
        status, data = self._req(
            "GET", "/tasks/%s/trace" % self.task_id, key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual([e["id"] for e in data["entries"]], ["c1"])
        status, data = self._req(
            "GET", "/tasks/%s/agents" % self.task_id, key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual([e["id"] for e in data["agents"][0]["calls"]],
                         ["c1"])

    def test_trace_merges_tenant_and_shared(self):
        self._write_tenant_trace([
            json.dumps({"kind": "exec", "id": "c1", "command": "tenant"}),
        ])
        self._write_trace([
            json.dumps({"kind": "exec", "id": "c2", "command": "shared"}),
        ])
        status, data = self._req(
            "GET", "/tasks/%s/trace" % self.task_id, key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual([e["id"] for e in data["entries"]], ["c1", "c2"])
        # 同 id 冲突时后读的共享目录覆盖（顺序有文档、可复现）。
        self._write_tenant_trace([
            json.dumps({"kind": "exec", "id": "c1",
                        "command": "tenant-cmd"}),
        ])
        self._write_trace([
            json.dumps({"kind": "exec", "id": "c1",
                        "command": "shared-cmd"}),
        ])
        status, data = self._req(
            "GET", "/tasks/%s/trace" % self.task_id, key=self.owner_key)
        self.assertEqual(data["entries"][0]["command"], "shared-cmd")

    def test_agents_tenant_child_file(self):
        # 委派子任务的子链路文件也在租户目录里：父标记的
        # trace 名要在租户目录里找得到，子任务调用才进链。
        self._write_tenant_trace([
            json.dumps({"kind": "agent-start", "id": "agent:rev",
                        "name": "rev", "prompt": "review",
                        "trace": "rev.jsonl", "ts": 1002}),
            json.dumps({"kind": "agent-done", "id": "agent:rev:done",
                        "name": "rev", "status": "completed",
                        "trace": "rev.jsonl", "ts": 1004}),
        ])
        self._write_tenant_trace([
            json.dumps({"kind": "exec", "id": "c9", "command": "grep"}),
        ], filename="rev.jsonl")
        status, data = self._req(
            "GET", "/tasks/%s/agents" % self.task_id, key=self.owner_key)
        self.assertEqual(status, 200)
        delegate = [a for a in data["agents"]
                    if a["kind"] == "delegate"][0]
        self.assertEqual(delegate["status"], "completed")
        self.assertEqual([e["id"] for e in delegate["calls"]], ["c9"])


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