#!/usr/bin/env python3
"""Tests for streamed Gitea sccache snapshot validation and merging."""

from __future__ import annotations

import hashlib
import importlib.util
import io
import json
from pathlib import Path
import os
import sys
import tarfile
import tempfile
import unittest
import zipfile


SCRIPT = Path(__file__).with_name("gitea_cache_snapshot.py")
SPEC = importlib.util.spec_from_file_location("gitea_cache_snapshot", SCRIPT)
assert SPEC and SPEC.loader
snapshot = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = snapshot
SPEC.loader.exec_module(snapshot)

REPOSITORY = "team/project"
SOURCE_SHA = "a" * 40
INHERITED_SOURCE_SHA = "c" * 40
BASE_IMAGE = "sha256:" + "b" * 64
RUSTC = "rustc 1.90.0\nbinary: rustc\n"
WORKSPACE = "/workspace/GenarrativeAI/Genarrative"
SCCACHE_VERSION = "0.18.0"


def object_path(character: str) -> str:
    key = character * 64
    return f"objects/{key[0]}/{key[1]}/{key}"


def cache_object(entries, *, mode=0o100644) -> bytes:
    output = io.BytesIO()
    with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_STORED) as bundle:
        for name, contents in entries:
            info = zipfile.ZipInfo(name)
            info.external_attr = mode << 16
            bundle.writestr(info, contents)
    return output.getvalue()


class SnapshotMergeTest(unittest.TestCase):
    def setUp(self) -> None:
        self.temporary_directory = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
        self.root = Path(self.temporary_directory.name).resolve()
        self.base = self.root / "base-objects"
        self.base.mkdir()

    def tearDown(self) -> None:
        self.temporary_directory.cleanup()

    def identity(self, job: str, run_id: int) -> object:
        return snapshot.ArtifactIdentity(REPOSITORY, run_id, 1, job, SOURCE_SHA)

    def archive(
        self,
        name: str,
        job: str,
        run_id: int,
        objects: list[tuple[str, bytes, int]],
        *,
        prefix: bool = True,
        manifest_update=None,
        extra_member: tuple[str, bytes, str] | None = None,
        touched: list[tuple[str, int]] | None = None,
    ) -> object:
        rows = [
            {
                "path": path,
                "size": len(contents),
                "sha256": hashlib.sha256(contents).hexdigest(),
                "mtime_ns": mtime_ns,
            }
            for path, contents, mtime_ns in objects
        ]
        manifest = {
            "schema": 1,
            "repository": REPOSITORY,
            "run_id": run_id,
            "run_attempt": 1,
            "job": job,
            "source_sha": SOURCE_SHA,
            "base_image": BASE_IMAGE,
            "rustc": RUSTC,
            "workspace": WORKSPACE,
            "sccache_version": SCCACHE_VERSION,
            "mode": "delta",
            "inherited_source_sha": INHERITED_SOURCE_SHA,
            "objects": rows,
            "touched": [
                {"path": path, "mtime_ns": mtime_ns}
                for path, mtime_ns in (touched or [])
            ],
        }
        if manifest_update is not None:
            manifest_update(manifest)

        tar_bytes = io.BytesIO()
        with tarfile.open(fileobj=tar_bytes, mode="w") as bundle:
            for path, contents, _ in objects:
                info = tarfile.TarInfo(path)
                info.size = len(contents)
                bundle.addfile(info, io.BytesIO(contents))
            if extra_member is not None:
                path, contents, kind = extra_member
                info = tarfile.TarInfo(path)
                if kind == "symlink":
                    info.type = tarfile.SYMTYPE
                    info.linkname = "manifest.json"
                    bundle.addfile(info)
                else:
                    info.size = len(contents)
                    bundle.addfile(info, io.BytesIO(contents))
            manifest_bytes = json.dumps(manifest, separators=(",", ":")).encode()
            info = tarfile.TarInfo("manifest.json")
            info.size = len(manifest_bytes)
            bundle.addfile(info, io.BytesIO(manifest_bytes))

        archive = self.root / f"{name}.zip"
        member = f"{name}/snapshot.tar" if prefix else "snapshot.tar"
        with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_STORED) as bundle:
            bundle.writestr(member, tar_bytes.getvalue())
        return snapshot.ArtifactInput(archive, self.identity(job, run_id))

    def base_object(self, path: str, contents: bytes, mtime_ns: int) -> None:
        target = self.base.joinpath(*Path(path).parts[1:])
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_bytes(contents)
        os.utime(target, ns=(mtime_ns, mtime_ns))

    def merge(self, inputs, output=None, maximum=1024 ** 3):
        return snapshot.merge_snapshots(
            inputs,
            output or self.root / "merged",
            base_objects=self.base,
            expected_inherited_source_sha=INHERITED_SOURCE_SHA,
            max_combined_bytes=maximum,
        )

    def test_merges_groups_deduplicates_newest_mtime_and_prunes_to_bound(self) -> None:
        first_path = object_path("a")
        second_path = object_path("b")
        third_path = object_path("c")
        inherited_path = object_path("d")
        self.base_object(inherited_path, b"dd", 50)
        inputs = [
            self.archive("first", "Backend tests", 41, [
                (first_path, b"aaaa", 100),
                (second_path, b"bbbb", 200),
            ]),
            self.archive("second", "Native shell tests", 41, [
                (first_path, b"aaaa", 300),
                (third_path, b"ccc", 250),
            ], prefix=False, touched=[(inherited_path, 400)]),
        ]

        output = self.root / "merged"
        result = self.merge(inputs, output, maximum=9)

        self.assertEqual((result.object_count, result.total_bytes), (3, 9))
        self.assertEqual((output / inherited_path).read_bytes(), b"dd")
        self.assertEqual((output / first_path).read_bytes(), b"aaaa")
        self.assertEqual((output / third_path).read_bytes(), b"ccc")
        self.assertFalse((output / second_path).exists())
        self.assertEqual((output / "rustc.txt").read_text(), RUSTC)
        self.assertEqual((output / "workspace.txt").read_text(), WORKSPACE + "\n")
        self.assertEqual((output / "source-commit.txt").read_text(), SOURCE_SHA + "\n")
        self.assertEqual((output / "base-image.txt").read_text(), BASE_IMAGE + "\n")
        self.assertFalse((output / "sccache").exists())

    def test_rejects_traversal_member_without_creating_output(self) -> None:
        item = self.archive(
            "traversal", "Backend tests", 42, [],
            extra_member=("../outside", b"bad", "file"),
        )
        output = self.root / "merged"
        with self.assertRaisesRegex(snapshot.SnapshotError, "invalid sccache object path"):
            self.merge([item], output)
        self.assertFalse(output.exists())
        self.assertFalse((self.root / "outside").exists())

    def test_rejects_checksum_mismatch(self) -> None:
        path = object_path("d")

        def corrupt(manifest):
            manifest["objects"][0]["sha256"] = "0" * 64

        item = self.archive(
            "checksum", "Backend tests", 43, [(path, b"content", 1)],
            manifest_update=corrupt,
        )
        with self.assertRaisesRegex(snapshot.SnapshotError, "checksums do not match"):
            self.merge([item])

    def test_rejects_manifest_identity_mismatch(self) -> None:
        item = self.archive(
            "identity", "Backend tests", 44, [],
            manifest_update=lambda manifest: manifest.update(run_id=999),
        )
        with self.assertRaisesRegex(snapshot.SnapshotError, "run_id does not match"):
            self.merge([item])

    def test_rejects_conflicting_duplicate_key(self) -> None:
        path = object_path("e")
        inputs = [
            self.archive("conflict-one", "Backend tests", 45, [(path, b"one", 1)]),
            self.archive("conflict-two", "Native shell tests", 45, [(path, b"two", 2)]),
        ]
        with self.assertRaisesRegex(snapshot.SnapshotError, "conflicting content"):
            self.merge(inputs)

    def test_deduplicates_cache_zip_member_order_preserving_payload_and_newest_touch(self) -> None:
        path = object_path("a")
        entries = [("lib.rlib", b"compiled"), ("lib.rmeta", b"metadata"), ("stderr", b"warning")]
        first = cache_object(entries)
        second = cache_object(list(reversed(entries)))
        self.assertEqual(len(first), len(second))
        self.assertNotEqual(hashlib.sha256(first).digest(), hashlib.sha256(second).digest())
        inputs = [
            self.archive("smoke", "smoke", 45, [(path, first, 100)]),
            self.archive("lane-1", "lane-1", 45, [(path, second, 300)]),
            self.archive("lane-2", "lane-2", 45, [(path, second, 200)]),
        ]
        result = self.merge(inputs)
        output = self.root / "merged" / path
        self.assertEqual((result.object_count, result.total_bytes), (1, len(first)))
        self.assertEqual(output.read_bytes(), first)
        self.assertEqual(output.stat().st_mtime_ns, 300)

    def test_rejects_cache_zip_payload_or_permissions_conflicts(self) -> None:
        path = object_path("a")
        first = cache_object([("lib.rlib", b"one"), ("stderr", b"err")])
        variants = {
            "payload": cache_object([("stderr", b"err"), ("lib.rlib", b"two")]),
            "permissions": cache_object([("stderr", b"err"), ("lib.rlib", b"one")], mode=0o100755),
        }
        for name, second in variants.items():
            with self.subTest(name=name):
                self.assertEqual(len(first), len(second))
                inputs = [
                    self.archive("first", "smoke", 45, [(path, first, 100)]),
                    self.archive("second", "lane-1", 45, [(path, second, 200)]),
                ]
                with self.assertRaisesRegex(snapshot.SnapshotError, "conflicting content"):
                    self.merge(inputs)
                self.assertFalse((self.root / "merged").exists())

    def test_rejects_delta_that_conflicts_with_inherited_key(self) -> None:
        path = object_path("f")
        self.base_object(path, b"base", 1)
        item = self.archive("base-conflict", "Backend tests", 46, [(path, b"evil", 2)])
        with self.assertRaisesRegex(snapshot.SnapshotError, "conflicting content"):
            self.merge([item])

    def test_accepts_empty_delta(self) -> None:
        item = self.archive("empty", "Backend tests", 47, [])
        output = self.root / "merged"
        result = self.merge([item], output)
        self.assertEqual((result.object_count, result.total_bytes), (0, 0))
        self.assertEqual(list((output / "objects").iterdir()), [])

    def test_rejects_noncanonical_base_entry(self) -> None:
        (self.base / "unexpected").write_bytes(b"not an object")
        item = self.archive("base-path", "Backend tests", 48, [])
        with self.assertRaisesRegex(snapshot.SnapshotError, "noncanonical entry"):
            self.merge([item])

    def test_rejects_tar_links(self) -> None:
        item = self.archive(
            "link", "Backend tests", 49, [],
            extra_member=(object_path("9"), b"", "symlink"),
        )
        with self.assertRaisesRegex(snapshot.SnapshotError, "regular file"):
            self.merge([item])


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