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_export,
    _publish_software_job_outputs,
    _publish_software_job_previews,
)


class SoftwareOutputPublishTests(unittest.TestCase):
    def test_workspace_preview_returns_to_task_folder_without_artifact_registration(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            job_id = uuid4()
            staging = root / ".zcbot_software_job_staging" / str(job_id)
            staging.mkdir(parents=True)
            content = b"preview"
            (staging / "figure.png").write_bytes(content)
            manifest = [{
                "artifact_id": "figure_png",
                "filename": "figure.png",
                "media_type": "image/png",
                "size_bytes": len(content),
                "sha256": hashlib.sha256(content).hexdigest(),
            }]
            context = {
                "capability": "origin.plot@v2",
                "user_id": uuid4(),
                "working_dir": str(root / "research"),
            }
            (root / "research").mkdir()
            with (
                patch("web.routers.software_nodes.load_user_root", return_value=root),
                patch("web.routers.software_nodes.register_published_artifacts") as register,
            ):
                previews = _publish_software_job_previews(job_id, context, manifest)
            self.assertFalse(register.called)
            self.assertNotIn("artifact_id", previews[0])
            self.assertEqual(previews[0]["preview_id"], "figure_png")
            self.assertEqual(previews[0]["path"], f"origin/{job_id}/figure.png")
            self.assertEqual(
                (root / "research" / "origin" / str(job_id) / "figure.png").read_bytes(),
                content,
            )

    def test_workspace_publishes_optional_previews_requested_by_the_job(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            job_id = uuid4()
            staging = root / ".zcbot_software_job_staging" / str(job_id)
            staging.mkdir(parents=True)
            files = {
                "preview_01": ("preview-01.png", b"overview"),
                "preview_02": ("preview-02.png", b"front"),
            }
            manifest = []
            for output_id, (filename, content) in files.items():
                (staging / filename).write_bytes(content)
                manifest.append({
                    "artifact_id": output_id,
                    "filename": filename,
                    "media_type": "image/png",
                    "size_bytes": len(content),
                    "sha256": hashlib.sha256(content).hexdigest(),
                })
            context = {
                "capability": "blender.scene.author@v3",
                "user_id": uuid4(),
                "working_dir": str(root / "research"),
                "request": {
                    "outputs": [
                        {"key": "preview_02", "type": "figure", "format": "png"}
                    ]
                },
            }
            (root / "research").mkdir()
            with patch("web.routers.software_nodes.load_user_root", return_value=root):
                previews = _publish_software_job_previews(job_id, context, manifest)
            self.assertEqual(
                {item["preview_id"] for item in previews},
                {"preview_01", "preview_02"},
            )

    def test_explicit_export_registers_an_existing_preview_file(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            job_id = uuid4()
            working_dir = root / "research"
            preview = working_dir / "origin" / str(job_id) / "figure.png"
            preview.parent.mkdir(parents=True)
            content = b"preview"
            preview.write_bytes(content)
            manifest = [{
                "artifact_id": "figure_png",
                "filename": "figure.png",
                "media_type": "image/png",
                "size_bytes": len(content),
                "sha256": hashlib.sha256(content).hexdigest(),
            }]
            context = {
                "capability": "origin.plot@v2",
                "user_id": uuid4(),
                "task_id": uuid4(),
                "working_dir": str(working_dir),
            }
            published_id = str(uuid4())
            with (
                patch("web.routers.software_nodes.load_user_root", return_value=root),
                patch(
                    "web.routers.software_nodes.register_published_artifacts",
                    return_value=({
                        "path": f"origin/{job_id}/figure.png",
                        "artifact_id": published_id,
                    },),
                ) as register,
            ):
                published = _publish_software_job_export(job_id, context, manifest)

            self.assertEqual(preview.read_bytes(), content)
            self.assertEqual(published[0]["artifact_id"], published_id)
            self.assertEqual(register.call_args.kwargs["software_job_id"], job_id)

    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()
