"""任务附件单测：上传包校验、建任务绑定、派发 seed、终局清理。

不跑真实模型回合、不联网：server 侧走本地临时端口 + 临时
SQLite；worker/sandbox 侧注入假 runner 录命令；host seed 侧用
临时目录。图片与普通文件不区分，一律是 worktree 里的文件。
"""

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

try:
    from cloud.auth import AuthStore
    from cloud.sandbox import DockerExecutor, LocalExecutor
    from cloud.server import create_server, init_db
    from cloud.uploads import (
        MAX_FILES,
        parse_multipart,
        parse_paths_field,
        save_bundle,
        validate_bundle,
        validate_rel_path,
    )
    from cloud.worker import Worker
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.sandbox import DockerExecutor, LocalExecutor
    from cloud.server import create_server, init_db
    from cloud.uploads import (
        MAX_FILES,
        parse_multipart,
        parse_paths_field,
        save_bundle,
        validate_bundle,
        validate_rel_path,
    )
    from cloud.worker import Worker


def _multipart_body(boundary, fields, files):
    """拼一个 multipart body：fields {名: 文本}，files
    [(part 名, 文件名, bytes)]。"""
    out = b""
    for name, value in fields.items():
        out += ("--%s\r\nContent-Disposition: form-data; name=\"%s\"\r\n\r\n"
                % (boundary, name)).encode("utf-8")
        out += value.encode("utf-8") + b"\r\n"
    for part_name, filename, data in files:
        out += ("--%s\r\nContent-Disposition: form-data; name=\"%s\"; "
                "filename=\"%s\"\r\nContent-Type: application/octet-stream"
                "\r\n\r\n" % (boundary, part_name, filename)).encode("utf-8")
        out += data + b"\r\n"
    out += ("--%s--\r\n" % boundary).encode("utf-8")
    return out


class UploadsModuleTestCase(unittest.TestCase):
    def test_parse_roundtrip_nested_paths(self):
        boundary = "TESTBOUNDARY"
        body = _multipart_body(
            boundary,
            {"paths": json.dumps(["a.txt", "sub/b.png"])},
            [("file0", "a.txt", b"hello"),
             ("file1", "b.png", b"\x89PNG")])
        fields, files = parse_multipart(
            body, "multipart/form-data; boundary=" + boundary)
        self.assertEqual(fields["paths"],
                         json.dumps(["a.txt", "sub/b.png"]))
        self.assertEqual([n for n, _d in files], ["a.txt", "b.png"])
        paths = parse_paths_field(fields, files)
        self.assertEqual(paths, ["a.txt", "sub/b.png"])

    def test_parse_paths_count_mismatch(self):
        with self.assertRaises(ValueError):
            parse_paths_field({"paths": json.dumps(["a.txt"])},
                              [("file0", b"x"), ("file1", b"y")])

    def test_validate_rel_path(self):
        for good in ["a.txt", "sub/b/c.png", "a-b_c.d/e f.txt", ".hidden/x"]:
            self.assertIsNone(validate_rel_path(good), good)
        for bad in ["/abs", "../x", "a/../b", "a//b", "", "a\\b",
                    "C:/x", "a\x00b", ".", "..", "a/./b"]:
            self.assertIsNotNone(validate_rel_path(bad), bad)

    def test_validate_bundle_limits(self):
        self.assertIsNotNone(validate_bundle(
            ["f%d.txt" % i for i in range(MAX_FILES + 1)],
            [1] * (MAX_FILES + 1)))
        self.assertIsNotNone(validate_bundle(
            ["big.bin"], [21 * 1024 * 1024]))
        self.assertIsNotNone(validate_bundle(
            ["a.bin", "b.bin"], [60 * 1024 * 1024, 60 * 1024 * 1024]))
        self.assertIsNone(validate_bundle(["a.txt"], [5]))

    def test_save_bundle_roundtrip_and_traversal(self):
        with tempfile.TemporaryDirectory() as tmp:
            bundle_id, saved = save_bundle(
                tmp, ["a.txt", "sub/b.txt"], [b"hello", b"world"])
            self.assertEqual([p for p, _s in saved],
                             ["a.txt", "sub/b.txt"])
            with open(os.path.join(tmp, bundle_id, "sub", "b.txt"),
                      "rb") as f:
                self.assertEqual(f.read(), b"world")
            with self.assertRaises(ValueError):
                save_bundle(tmp, ["../escape.txt"], [b"x"])


class ServerAttachmentsTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")
        self.upload_dir = os.path.join(self.tmp.name, "uploads")
        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.upload_dir = self.upload_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,
             content_type="application/json"):
        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:
            if isinstance(body, bytes):
                payload = body
            else:
                payload = json.dumps(body)
            hdrs["Content-Type"] = content_type
        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 _upload(self, paths, contents, key=None):
        body = _multipart_body(
            "TESTBOUNDARY", {"paths": json.dumps(paths)},
            [("file%d" % i, os.path.basename(p) or "f", data)
             for i, (p, data) in enumerate(zip(paths, contents))])
        return self._req("POST", "/uploads", body,
                         key or self.owner_key,
                         content_type="multipart/form-data;"
                         " boundary=TESTBOUNDARY")

    def _create_task(self, upload_id=None, executor=None):
        body = {"project_id": self.project["id"], "prompt": "do work"}
        if upload_id is not None:
            body["upload_id"] = upload_id
        if executor is not None:
            body["executor"] = executor
        return self._req("POST", "/tasks", body, self.owner_key)

    def test_upload_and_bind(self):
        status, up = self._upload(["a.txt", "sub/b.png"],
                                  [b"hello", b"\x89PNG"])
        self.assertEqual(status, 201, up)
        self.assertEqual([f["path"] for f in up["files"]],
                         ["a.txt", "sub/b.png"])
        status, task = self._create_task(up["upload_id"])
        self.assertEqual(status, 201, task)
        self.assertEqual(task["attachments"], ["a.txt", "sub/b.png"])
        status, got = self._req(
            "GET", "/tasks/" + task["id"], key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(got["attachments"], ["a.txt", "sub/b.png"])
        status, listed = self._req(
            "GET", "/tasks?project_id=" + self.project["id"],
            key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(listed["tasks"][0]["attachments"],
                         ["a.txt", "sub/b.png"])

    def test_upload_rejects_bad_path_and_empty(self):
        status, body = self._upload(["../escape.txt"], [b"x"])
        self.assertEqual(status, 400)
        status, body = self._req("POST", "/uploads",
                                 {"paths": "[]"}, self.owner_key)
        self.assertEqual(status, 400)

    def test_bind_rejects_used_unknown_other_cloud(self):
        status, up = self._upload(["a.txt"], [b"hi"])
        self.assertEqual(status, 201)
        status, _task = self._create_task(up["upload_id"])
        self.assertEqual(status, 201)
        status, body = self._create_task(up["upload_id"])
        self.assertEqual(status, 409)
        status, body = self._create_task("no-such-bundle")
        self.assertEqual(status, 400)
        status, up2 = self._upload(["b.txt"], [b"hi"])
        self.assertEqual(status, 201)
        status, body = self._req(
            "POST", "/tasks",
            {"project_id": self.project["id"], "prompt": "x",
             "upload_id": up2["upload_id"]}, self.other_key)
        self.assertEqual(status, 403)
        status, up3 = self._upload(["c.txt"], [b"hi"])
        self.assertEqual(status, 201)
        status, body = self._create_task(up3["upload_id"],
                                         executor="cloud")
        self.assertEqual(status, 400)

    def test_cancel_deletes_bundle_dir(self):
        status, up = self._upload(["a.txt"], [b"hi"])
        self.assertEqual(status, 201)
        bundle_dir = os.path.join(self.upload_dir, up["upload_id"])
        self.assertTrue(os.path.isdir(bundle_dir))
        status, task = self._create_task(up["upload_id"])
        self.assertEqual(status, 201)
        status, _body = self._req(
            "POST", "/tasks/%s/cancel" % task["id"], {},
            self.owner_key)
        self.assertEqual(status, 200, _body)
        self.assertFalse(os.path.exists(bundle_dir))

    def test_fork_copies_attachment_rows(self):
        status, up = self._upload(["a.txt"], [b"hi"])
        self.assertEqual(status, 201)
        status, task = self._create_task(up["upload_id"])
        self.assertEqual(status, 201)
        conn = sqlite3.connect(self.db_path)
        conn.execute("UPDATE tasks SET status = 'completed' WHERE id = ?",
                     (task["id"],))
        conn.commit()
        conn.close()
        status, forked = self._req(
            "POST", "/tasks/%s/fork" % task["id"],
            {"message": "continue"}, self.owner_key)
        self.assertEqual(status, 201, forked)
        self.assertEqual(forked["attachments"], ["a.txt"])


class SandboxSeedTestCase(unittest.TestCase):
    def _recording(self, calls):
        def fake(cmd):
            calls.append(list(cmd))
            return SimpleNamespace(returncode=0, stdout="ok", stderr="")
        return fake

    def test_local_dispatch_and_goal_seed_flags(self):
        calls = []
        ex = LocalExecutor("/fake/host.py", runner=self._recording(calls))
        ex.dispatch_task({"id": "t1", "prompt": "do"},
                         seed_dir="/tmp/seed")
        self.assertIn("--seed-dir", calls[0])
        self.assertEqual(calls[0][calls[0].index("--seed-dir") + 1],
                         "/tmp/seed")
        calls.clear()
        ex.dispatch_task({"id": "t1", "prompt": "do"})
        self.assertNotIn("--seed-dir", calls[0])
        calls.clear()
        ex.goal_task({"id": "t2", "prompt": "build it", "goal_mode": 1,
                      "goal": "build it", "goal_max_rounds": 0},
                     seed_dir="/tmp/seed")
        self.assertIn("--seed-dir", calls[0])
        self.assertNotIn("--verify", calls[0])
        self.assertEqual(calls[0][calls[0].index("--max-rounds") + 1],
                         "0")

    def test_popen_goal_new_semantics_and_seed(self):
        import cloud.sandbox as sandbox_mod
        calls = []

        class FakePopen(object):
            def __init__(self, cmd, **kwargs):
                calls.append(list(cmd))

        orig = sandbox_mod.subprocess.Popen
        sandbox_mod.subprocess.Popen = FakePopen
        try:
            ex = LocalExecutor("/fake/host.py")
            ex.popen_task({"id": "t2", "prompt": "build it",
                           "goal_mode": 1, "goal": "build it",
                           "goal_max_rounds": 0},
                          kind="goal", seed_dir="/tmp/seed")
            ex.popen_task({"id": "t1", "prompt": "do"}, kind="run",
                          seed_dir="/tmp/seed")
        finally:
            sandbox_mod.subprocess.Popen = orig
        goal_cmd, run_cmd = calls
        self.assertNotIn("--verify", goal_cmd)
        self.assertEqual(goal_cmd[goal_cmd.index("--max-rounds") + 1],
                         "0")
        self.assertIn("--seed-dir", goal_cmd)
        self.assertNotIn("--seed-dir", run_cmd)

    def test_docker_seed_mount(self):
        from unittest import mock
        calls = []
        ex = DockerExecutor(
            "/fake/host.py", runner=self._recording(calls),
            env={"XCM_API_KEY": "placeholder-not-a-real-key"})
        with tempfile.TemporaryDirectory() as seed_dir, \
                mock.patch("cloud.sandbox.shutil.which",
                           return_value="/usr/bin/docker"):
            ex.dispatch_task({"id": "t1", "prompt": "do"},
                             seed_host_dir=seed_dir)
            flat = " ".join(calls[0])
            self.assertIn("%s:/tmp/codex-seed:ro" % seed_dir, flat)
            self.assertIn("/tmp/codex-seed", flat)
            calls.clear()
            ex.dispatch_task({"id": "t1", "prompt": "do"})
            flat = " ".join(calls[0])
            self.assertNotIn("codex-seed", flat)


class WorkerSeedTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")
        self.upload_dir = os.path.join(self.tmp.name, "uploads")
        init_db(self.db_path)
        self.env = {"XCM_API_KEY": "test-placeholder"}

    def tearDown(self):
        self.tmp.cleanup()

    def _insert_task(self, task_id, **extra):
        now = time.time()
        cols = ["id", "project_id", "prompt", "status",
                "idempotency_key", "depends_on", "priority",
                "created_by", "created_at", "updated_at",
                "last_heartbeat"]
        vals = [task_id, "p1", "do work", "queued", None, "[]", 0,
                "tester", now, now, None]
        for k, v in extra.items():
            cols.append(k)
            vals.append(v)
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute(
                "INSERT INTO tasks(%s) VALUES(%s)"
                % (", ".join(cols), ", ".join("?" * len(cols))), vals)
            conn.commit()
        finally:
            conn.close()

    def _insert_bundle(self, task_id, bundle_id):
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute(
                "INSERT INTO task_attachments"
                "(task_id, bundle_id, rel_path, size) VALUES(?, ?, ?, ?)",
                (task_id, bundle_id, "a.txt", 2))
            conn.commit()
        finally:
            conn.close()

    def test_task_seed_dir_resolution(self):
        worker = Worker(self.db_path, "/fake/host.py", env=self.env,
                        upload_dir=self.upload_dir)
        self.assertIsNone(worker._task_seed_dir({"id": "nope"}))
        self._insert_bundle("t1", "b1")
        self.assertIsNone(worker._task_seed_dir({"id": "t1"}))
        bundle_dir = os.path.join(self.upload_dir, "b1")
        os.makedirs(bundle_dir)
        self.assertEqual(worker._task_seed_dir({"id": "t1"}), bundle_dir)

    def test_run_one_passes_seed_and_cleans_bundle(self):
        self._insert_task("t1")
        self._insert_bundle("t1", "b1")
        bundle_dir = os.path.join(self.upload_dir, "b1")
        os.makedirs(bundle_dir)
        with open(os.path.join(bundle_dir, "a.txt"), "w") as f:
            f.write("hi")
        seen = []

        def fake(cmd):
            seen.append(list(cmd))
            return SimpleNamespace(returncode=0, stdout="ok", stderr="")

        worker = Worker(self.db_path, "/fake/host.py", runner=fake,
                        env=self.env, upload_dir=self.upload_dir)
        result = worker.run_one()
        self.assertEqual(result["status"], "completed")
        dispatch_cmds = [c for c in seen if "dispatch" in c]
        self.assertTrue(dispatch_cmds)
        self.assertIn("--seed-dir", dispatch_cmds[0])
        self.assertEqual(
            dispatch_cmds[0][dispatch_cmds[0].index("--seed-dir") + 1],
            bundle_dir)
        self.assertFalse(os.path.exists(bundle_dir))


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