from __future__ import annotations

import hashlib
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
from uuid import uuid4

from web.routers.software_nodes import _publish_software_job_outputs


class SoftwareOutputPublishTests(unittest.TestCase):
    def test_complete_set_moves_atomically_and_can_be_replayed(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            job_id = uuid4()
            working_dir = root / "research"
            staging = root / ".zcbot_software_job_staging" / str(job_id)
            staging.mkdir(parents=True)
            working_dir.mkdir()
            content = b"origin-result"
            metadata = b'{"kind":"plot"}'
            (staging / "figure.png").write_bytes(content)
            (staging / "plot-spec.json").write_bytes(metadata)
            manifest = [
                {
                    "artifact_id": "figure_png",
                    "filename": "figure.png",
                    "media_type": "image/png",
                    "size_bytes": len(content),
                    "sha256": hashlib.sha256(content).hexdigest(),
                },
                {
                    "artifact_id": "plot_spec",
                    "filename": "plot-spec.json",
                    "media_type": "application/json",
                    "size_bytes": len(metadata),
                    "sha256": hashlib.sha256(metadata).hexdigest(),
                },
            ]
            context = {
                "user_id": uuid4(),
                "task_id": uuid4(),
                "working_dir": str(working_dir),
            }

            registered = []

            def register(**kwargs):
                registered.append(kwargs)
                return tuple({
                    "version": 2,
                    "scope": "working_dir",
                    "path": ref["path"],
                    "label": ref["label"],
                    "artifact_id": str(uuid4()),
                } for ref in kwargs["refs"])

            with (
                patch("web.routers.software_nodes.load_user_root", return_value=root),
                patch(
                    "web.routers.software_nodes.register_published_artifacts",
                    side_effect=register,
                ),
            ):
                first = _publish_software_job_outputs(job_id, context, manifest)
                second = _publish_software_job_outputs(job_id, context, manifest)

            published = working_dir / "origin" / str(job_id) / "figure.png"
            self.assertEqual(published.read_bytes(), content)
            self.assertEqual(
                (working_dir / "origin" / str(job_id) / ".meta" / "plot-spec.json").read_bytes(),
                metadata,
            )
            self.assertFalse(staging.exists())
            self.assertEqual(first[0]["source_artifact_id"], "figure_png")
            self.assertEqual(first[0]["path"], f"origin/{job_id}/figure.png")
            self.assertEqual(first[1]["source_artifact_id"], "plot_spec")
            self.assertIsNone(first[1]["artifact_id"])
            self.assertEqual(first[1]["path"], f"origin/{job_id}/.meta/plot-spec.json")
            self.assertEqual(second[0]["source_artifact_id"], "figure_png")
            self.assertEqual(registered[0]["software_job_id"], job_id)
            self.assertEqual(len(registered[0]["refs"]), 1)


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