"""Unit: the infra-artifact publisher's upload path (PRD 0069 Stage 2).

The rootfs is hundreds of MB, so `_put` must stream it from disk rather than
read it into memory. Network is mocked; no Docker, no real build.
"""

from __future__ import annotations

import hashlib
from email.message import Message
import tempfile
import unittest
import urllib.error
import urllib.request
from pathlib import Path
from unittest import mock

from bot_bottle.backend.firecracker import publish_infra as pub


class _Resp:
    status = 201

    def __enter__(self) -> "_Resp":
        return self

    def __exit__(self, *a: object) -> bool:
        return False


class TestPut(unittest.TestCase):
    def test_streams_file_body_with_content_length(self) -> None:
        captured: list[urllib.request.Request] = []

        def fake_urlopen(req: urllib.request.Request, *a: object, **k: object) -> _Resp:
            captured.append(req)
            return _Resp()

        with tempfile.TemporaryDirectory() as d:
            f = Path(d) / "rootfs.ext4.gz"
            payload = b"x" * 4096
            f.write_bytes(payload)
            with mock.patch.object(pub.urllib.request, "urlopen", fake_urlopen):
                pub._put("https://reg/pkg", f, token="t")

        req = captured[0]
        # Body is the open file object (streamed), never the bytes in memory.
        self.assertTrue(hasattr(req.data, "read"))
        self.assertNotIsInstance(req.data, (bytes, bytearray))
        self.assertEqual(str(len(payload)), req.get_header("Content-length"))

    def test_small_bytes_body_still_works(self) -> None:
        captured: list[urllib.request.Request] = []

        def fake_urlopen(req: urllib.request.Request, *a: object, **k: object) -> _Resp:
            captured.append(req)
            return _Resp()

        with mock.patch.object(pub.urllib.request, "urlopen", fake_urlopen):
            pub._put("https://reg/sha", b"abc123  rootfs\n", token="")
        self.assertEqual(b"abc123  rootfs\n", captured[0].data)


_ROLE = "orchestrator"


class TestPublishBundle(unittest.TestCase):
    def _bundle(self, role_dir: Path, version: str) -> None:
        payload = b"candidate"
        role_dir.mkdir(parents=True, exist_ok=True)
        (role_dir / "version.txt").write_text(version + "\n")
        (role_dir / "rootfs.ext4.gz").write_bytes(payload)
        digest = hashlib.sha256(payload).hexdigest()
        (role_dir / "rootfs.ext4.gz.sha256").write_text(
            f"{digest}  rootfs.ext4.gz\n")

    def test_existing_identical_artifact_is_success(self) -> None:
        with tempfile.TemporaryDirectory() as d:
            role_dir = Path(d) / _ROLE
            self._bundle(role_dir, "v1")
            sha = (role_dir / "rootfs.ext4.gz.sha256").read_bytes()
            response = mock.MagicMock()
            response.__enter__.return_value.read.return_value = sha
            with mock.patch.object(pub, "_role_version", return_value="v1"), \
                    mock.patch.object(
                        pub.urllib.request, "urlopen", return_value=response
                    ), mock.patch.object(pub, "_put") as put:
                self.assertEqual("v1", pub._publish_bundle(_ROLE, role_dir, "token"))
            put.assert_not_called()

    def test_partial_artifact_is_replaced(self) -> None:
        with tempfile.TemporaryDirectory() as d:
            role_dir = Path(d) / _ROLE
            self._bundle(role_dir, "v1")
            missing = urllib.error.HTTPError("u", 404, "missing", Message(), None)
            with mock.patch.object(pub, "_role_version", return_value="v1"), \
                    mock.patch.object(
                        pub.urllib.request, "urlopen", side_effect=missing
                    ), mock.patch.object(pub, "_delete") as delete, \
                    mock.patch.object(pub, "_put") as put:
                pub._publish_bundle(_ROLE, role_dir, "token")
            self.assertEqual(3, delete.call_count)
            self.assertEqual(3, put.call_count)

    def test_rejects_bundle_for_different_checkout(self) -> None:
        with tempfile.TemporaryDirectory() as d:
            role_dir = Path(d) / _ROLE
            self._bundle(role_dir, "old")
            with mock.patch.object(pub, "_role_version", return_value="new"):
                with self.assertRaises(SystemExit) as ctx:
                    pub._publish_bundle(_ROLE, role_dir, "token")
            self.assertIn("does not match checkout", str(ctx.exception))

    def test_rejects_bad_bundle_checksum(self) -> None:
        with tempfile.TemporaryDirectory() as d:
            role_dir = Path(d) / _ROLE
            self._bundle(role_dir, "v1")
            (role_dir / "rootfs.ext4.gz").write_bytes(b"tampered")
            with mock.patch.object(pub, "_role_version", return_value="v1"):
                with self.assertRaises(SystemExit) as ctx:
                    pub._publish_bundle(_ROLE, role_dir, "token")
            self.assertIn("checksum mismatch", str(ctx.exception))

    def test_registry_lookup_failure_is_reported(self) -> None:
        with tempfile.TemporaryDirectory() as d:
            role_dir = Path(d) / _ROLE
            self._bundle(role_dir, "v1")
            failure = urllib.error.URLError("offline")
            with mock.patch.object(pub, "_role_version", return_value="v1"), \
                    mock.patch.object(pub.urllib.request, "urlopen", side_effect=failure):
                with self.assertRaises(SystemExit) as ctx:
                    pub._publish_bundle(_ROLE, role_dir, "token")
            self.assertIn("registry unreachable", str(ctx.exception))


class TestTryDownloadPublished(unittest.TestCase):
    def test_downloads_existing_artifact(self) -> None:
        with tempfile.TemporaryDirectory() as d, \
                mock.patch.object(pub, "_role_version", return_value="v1"), \
                mock.patch.object(
                    pub.urllib.request, "urlopen", return_value=_Resp()
                ), mock.patch.object(pub.infra_artifact, "_download") as download:
            role_dir = Path(d) / _ROLE
            role_dir.mkdir()
            result = pub._try_download_published(_ROLE, role_dir)
            self.assertEqual("v1", result)
            self.assertEqual("v1\n", (role_dir / "version.txt").read_text())
            self.assertEqual(2, download.call_count)

    def test_missing_artifact_returns_none(self) -> None:
        missing = urllib.error.HTTPError("u", 404, "missing", Message(), None)
        with tempfile.TemporaryDirectory() as d, \
                mock.patch.object(pub, "_role_version", return_value="v1"), \
                mock.patch.object(pub.urllib.request, "urlopen", side_effect=missing):
            self.assertIsNone(pub._try_download_published(_ROLE, Path(d)))

    def test_registry_http_failure_is_reported(self) -> None:
        failure = urllib.error.HTTPError("u", 500, "failed", Message(), None)
        with tempfile.TemporaryDirectory() as d, \
                mock.patch.object(pub, "_role_version", return_value="v1"), \
                mock.patch.object(pub.urllib.request, "urlopen", side_effect=failure):
            with self.assertRaises(SystemExit) as ctx:
                pub._try_download_published(_ROLE, Path(d))
        self.assertIn("registry check failed (HTTP 500)", str(ctx.exception))

    def test_registry_connection_failure_is_reported(self) -> None:
        with tempfile.TemporaryDirectory() as d, \
                mock.patch.object(pub, "_role_version", return_value="v1"), \
                mock.patch.object(
                    pub.urllib.request, "urlopen",
                    side_effect=urllib.error.URLError("offline")):
            with self.assertRaises(SystemExit) as ctx:
                pub._try_download_published(_ROLE, Path(d))
        self.assertIn("registry unreachable", str(ctx.exception))


class TestMain(unittest.TestCase):
    def test_output_builds_a_candidate_per_role(self) -> None:
        with tempfile.TemporaryDirectory() as d:
            root = Path(d) / "candidate"
            with mock.patch.object(pub.infra_vm, "build_infra_images_with_docker") as images, \
                    mock.patch.object(pub, "build_role_artifact", return_value="v1") as build:
                self.assertEqual(0, pub.main(["--output", str(root)]))
            images.assert_called_once()
            built_roles = {c.args[0] for c in build.call_args_list}
            self.assertEqual({"orchestrator", "gateway"}, built_roles)

    def test_output_reuses_published_candidates(self) -> None:
        with tempfile.TemporaryDirectory() as d:
            root = Path(d) / "candidate"
            with mock.patch.object(pub, "_try_download_published", return_value="v1") as reuse, \
                    mock.patch.object(pub.infra_vm, "build_infra_images_with_docker") as images, \
                    mock.patch.object(pub, "build_role_artifact") as build:
                self.assertEqual(
                    0, pub.main(["--output", str(root), "--reuse-published"]))
            self.assertEqual(2, reuse.call_count)   # once per role
            images.assert_not_called()
            build.assert_not_called()

    def test_publish_dir_publishes_both_roles(self) -> None:
        with tempfile.TemporaryDirectory() as d, \
                mock.patch.object(pub.infra_artifact, "_config", return_value=("", "", "t")), \
                mock.patch.object(pub, "_publish_bundle", return_value="v1") as publish:
            self.assertEqual(0, pub.main(["--publish-dir", d]))
            published_roles = {c.args[0] for c in publish.call_args_list}
            self.assertEqual({"orchestrator", "gateway"}, published_roles)


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