from __future__ import annotations

import json
import stat
import tempfile
import unittest
from dataclasses import replace
from pathlib import Path

from uvm.config import Settings
from uvm.domain import VmRecord
from uvm.state import StateStore


def make_vm(vm_id: str = "vm-test") -> VmRecord:
    return VmRecord(
        id=vm_id,
        cpu=1,
        ram_mib=512,
        guest_ip="10.42.0.2",
        gateway="10.42.0.1",
        tap="uvm-test",
        mac="02:fc:00:00:00:01",
        socket="/tmp/firecracker.sock",
        config="/tmp/config.json",
        log="/tmp/firecracker.log",
        disk="/tmp/rootfs.ext4",
        username="admin",
        password="stored-secret",
    )


class StateStoreTests(unittest.TestCase):
    def setUp(self) -> None:
        self._temporary_directory = tempfile.TemporaryDirectory()
        self.base = Path(self._temporary_directory.name) / "uvm"
        self.settings = replace(Settings(), base=self.base)
        self.store = StateStore(self.settings)

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

    def test_transaction_persists_a_vm_atomically(self) -> None:
        self.store.initialize()
        with self.store.transaction() as state:
            state.vms["vm-test"] = make_vm()
            state.next_mac_index = 2

        loaded = self.store.load()
        self.assertEqual(loaded.vms["vm-test"].disk, "/tmp/rootfs.ext4")
        self.assertEqual(loaded.vms["vm-test"].username, "admin")
        self.assertEqual(loaded.vms["vm-test"].password, "stored-secret")
        self.assertNotIn("stored-secret", repr(loaded.vms["vm-test"]))
        self.assertEqual(loaded.next_mac_index, 2)
        self.assertEqual(stat.S_IMODE(self.settings.state_path.stat().st_mode), 0o600)
        persisted = json.loads(self.settings.state_path.read_text(encoding="utf-8"))
        self.assertEqual(persisted["vms"]["vm-test"]["username"], "admin")
        self.assertEqual(persisted["vms"]["vm-test"]["password"], "stored-secret")

    def test_legacy_state_is_loaded_and_migrated_on_next_write(self) -> None:
        self.base.mkdir(parents=True)
        legacy = {"vms": {"vm-test": make_vm().to_dict()}}
        legacy["vms"]["vm-test"].pop("disk")
        legacy["vms"]["vm-test"].pop("username")
        legacy["vms"]["vm-test"].pop("password")
        legacy["vms"]["vm-test"].pop("updated_at")
        self.settings.state_path.write_text(json.dumps(legacy), encoding="utf-8")

        loaded = self.store.load()
        self.assertEqual(loaded.vms["vm-test"].disk, "")
        self.assertEqual(loaded.vms["vm-test"].username, "root")
        self.assertIsNone(loaded.vms["vm-test"].password)
        self.assertEqual(loaded.next_mac_index, 2)

        with self.store.transaction() as state:
            state.vms["vm-test"].status = "stopped"

        persisted = json.loads(self.settings.state_path.read_text(encoding="utf-8"))
        self.assertEqual(persisted["schema_version"], 1)
        self.assertEqual(persisted["next_mac_index"], 2)
        self.assertEqual(persisted["vms"]["vm-test"]["status"], "stopped")

    def test_initialize_protects_registry_credentials(self) -> None:
        self.base.mkdir(parents=True)
        self.settings.state_path.write_text('{"vms": {}}', encoding="utf-8")
        self.settings.state_path.chmod(0o644)

        self.store.initialize()

        self.assertEqual(stat.S_IMODE(self.settings.state_path.stat().st_mode), 0o600)