"""任务列表 total_tokens：一次聚合，不逐个调成本接口。

用真实 HTTP 服务 + 临时 SQLite：建两个任务，给其中一个记
两条成本记录（actual + estimated），列表应一次回两行的
total_tokens（有记录的加总、无记录的 0），且不回 404/500；
_task_tokens_map 在 cost_records 表不存在时也不炸列表。
"""

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

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


class ListTokensTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")
        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.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 body is not None:
            payload = json.dumps(body).encode("utf-8")
            hdrs["Content-Type"] = "application/json"
        if key is not None:
            hdrs["X-API-Key"] = key
        try:
            conn.request(method, path, body=payload, headers=hdrs)
            resp = conn.getresponse()
            data = resp.read()
        finally:
            conn.close()
        try:
            parsed = json.loads(data.decode("utf-8")) if data else None
        except ValueError:
            parsed = None
        return resp.status, parsed

    def test_list_returns_aggregated_total_tokens(self):
        pid = self.project["id"]
        _, t1 = self._req("POST", "/tasks",
                          {"project_id": pid, "prompt": "有成本"},
                          self.owner_key)
        _, t2 = self._req("POST", "/tasks",
                          {"project_id": pid, "prompt": "无成本"},
                          self.owner_key)
        store = CostStore(self.db_path)
        try:
            store.record(t1["id"], pid, "m", input_tokens=100,
                         output_tokens=50, source="actual")
            store.record(t1["id"], pid, "m", input_tokens=20,
                         output_tokens=0, source="estimated")
        finally:
            store._conn.close()
        status, body = self._req(
            "GET", "/tasks?project_id=%s" % pid, key=self.owner_key)
        self.assertEqual(status, 200, body)
        by_id = {t["id"]: t for t in body["tasks"]}
        self.assertEqual(by_id[t1["id"]]["total_tokens"], 170)
        self.assertEqual(by_id[t2["id"]]["total_tokens"], 0)

    def test_single_task_returns_aggregated_total_tokens(self):
        # 详情接口（_get_task）与列表口径一致：单任务也聚合回
        # total_tokens，否则详情里的 token 栏只能显示占位。
        pid = self.project["id"]
        _, t1 = self._req("POST", "/tasks",
                          {"project_id": pid, "prompt": "详情成本"},
                          self.owner_key)
        store = CostStore(self.db_path)
        try:
            store.record(t1["id"], pid, "m", input_tokens=60,
                         output_tokens=40, source="actual")
        finally:
            store._conn.close()
        status, body = self._req(
            "GET", "/tasks/%s" % t1["id"], key=self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["total_tokens"], 100)

    def test_tokens_map_missing_table_is_zero_not_error(self):
        import sqlite3
        conn = sqlite3.connect(":memory:")
        try:
            self.assertEqual(
                _task_tokens_map(conn, ["t1", "t2"]),
                {"t1": 0, "t2": 0})
        finally:
            conn.close()

    def test_task_to_dict_default_total_tokens_is_none(self):
        row = {"id": "t1", "project_id": "p1", "prompt": "p",
               "status": "completed", "idempotency_key": None,
               "depends_on": "[]", "priority": 0, "created_by": "u",
               "created_at": 1.0, "updated_at": 2.0,
               "last_heartbeat": None}
        self.assertIsNone(_task_to_dict(row, [])["total_tokens"])
        self.assertEqual(
            _task_to_dict(row, [], 42)["total_tokens"], 42)


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