from datetime import UTC, datetime

import pytest
from pydantic import ValidationError

from atlas_runner.models import (
    ContainerState,
    ContainerStatus,
    Mount,
    RunResult,
    RunSpec,
)


class TestMount:
    def test_defaults(self) -> None:
        m = Mount(src="/host", dst="/container")
        assert m.src == "/host"
        assert m.dst == "/container"
        assert m.ro is False

    def test_readonly(self) -> None:
        m = Mount(src="/data", dst="/mnt", ro=True)
        assert m.ro is True

    def test_round_trip(self) -> None:
        m = Mount(src="/a", dst="/b", ro=True)
        data = m.model_dump()
        assert data == {"src": "/a", "dst": "/b", "ro": True}
        m2 = Mount.model_validate(data)
        assert m2 == m


class TestContainerState:
    def test_all_values(self) -> None:
        assert ContainerState.RUNNING == "running"
        assert ContainerState.EXITED == "exited"
        assert ContainerState.DEAD == "dead"
        assert ContainerState.CREATED == "created"
        assert ContainerState.PAUSED == "paused"
        assert ContainerState.RESTARTING == "restarting"
        assert ContainerState.REMOVING == "removing"

    def test_member_count(self) -> None:
        assert len(ContainerState) == 7


class TestRunSpec:
    def test_minimal_with_command(self) -> None:
        spec = RunSpec(name="test", image="alpine", command=["echo", "hi"])
        assert spec.name == "test"
        assert spec.image == "alpine"
        assert spec.command == ["echo", "hi"]
        assert spec.shell is None

    def test_minimal_with_shell(self) -> None:
        spec = RunSpec(name="test", image="alpine", shell="echo hi")
        assert spec.shell == "echo hi"
        assert spec.command is None

    def test_no_command_no_shell(self) -> None:
        spec = RunSpec(name="test", image="alpine")
        assert spec.command is None
        assert spec.shell is None

    def test_command_and_shell_rejected(self) -> None:
        with pytest.raises(ValidationError, match="mutually exclusive"):
            RunSpec(name="test", image="alpine", command=["echo"], shell="echo")

    def test_cpu_cores_only(self) -> None:
        spec = RunSpec(name="test", image="alpine", cpu_cores=4)
        assert spec.cpu_cores == 4
        assert spec.cpuset is None

    def test_cpuset_only(self) -> None:
        spec = RunSpec(name="test", image="alpine", cpuset="0-3")
        assert spec.cpuset == "0-3"
        assert spec.cpu_cores is None

    def test_cpu_cores_and_cpuset_rejected(self) -> None:
        with pytest.raises(ValidationError, match="mutually exclusive"):
            RunSpec(name="test", image="alpine", cpu_cores=4, cpuset="0-3")

    def test_defaults(self) -> None:
        spec = RunSpec(name="test", image="alpine")
        assert spec.labels == {}
        assert spec.gpus is None
        assert spec.memory_gb is None
        assert spec.scratch is None
        assert spec.mounts == []
        assert spec.network is None
        assert spec.ports == {}
        assert spec.workdir is None
        assert spec.env == {}
        assert spec.cap_drop_all is True
        assert spec.no_new_privileges is True

    def test_full_spec(self) -> None:
        spec = RunSpec(
            name="full",
            image="nvidia/cuda:12.0-base",
            command=["python", "train.py"],
            labels={"project": "atlas"},
            gpus=[0, 1],
            cpu_cores=8,
            memory_gb=16,
            scratch="/tmp/scratch",
            mounts=[Mount(src="/data", dst="/mnt/data", ro=True)],
            network="host",
            ports={"8080": 80},
            workdir="/app",
            env={"KEY": "value"},
            cap_drop_all=False,
            no_new_privileges=False,
        )
        assert spec.gpus == [0, 1]
        assert spec.mounts[0].ro is True
        assert spec.env["KEY"] == "value"

    def test_round_trip(self) -> None:
        spec = RunSpec(name="rt", image="alpine", command=["ls"])
        data = spec.model_dump()
        spec2 = RunSpec.model_validate(data)
        assert spec2 == spec


class TestRunResult:
    def test_fields(self) -> None:
        ts = datetime(2026, 1, 1, tzinfo=UTC)
        r = RunResult(container_id="abc123", name="test", started_at=ts)
        assert r.container_id == "abc123"
        assert r.name == "test"
        assert r.started_at == ts

    def test_round_trip(self) -> None:
        ts = datetime(2026, 1, 1, tzinfo=UTC)
        r = RunResult(container_id="abc", name="t", started_at=ts)
        data = r.model_dump()
        r2 = RunResult.model_validate(data)
        assert r2 == r


class TestContainerStatus:
    def test_running(self) -> None:
        ts = datetime(2026, 1, 1, tzinfo=UTC)
        s = ContainerStatus(
            container_id="abc",
            name="test",
            state=ContainerState.RUNNING,
            started_at=ts,
        )
        assert s.state == ContainerState.RUNNING
        assert s.exit_code is None
        assert s.finished_at is None

    def test_exited(self) -> None:
        ts_start = datetime(2026, 1, 1, tzinfo=UTC)
        ts_end = datetime(2026, 1, 1, 0, 5, tzinfo=UTC)
        s = ContainerStatus(
            container_id="abc",
            name="test",
            state=ContainerState.EXITED,
            exit_code=0,
            started_at=ts_start,
            finished_at=ts_end,
        )
        assert s.exit_code == 0
        assert s.finished_at == ts_end

    def test_round_trip(self) -> None:
        s = ContainerStatus(
            container_id="x",
            name="y",
            state=ContainerState.CREATED,
        )
        data = s.model_dump()
        s2 = ContainerStatus.model_validate(data)
        assert s2 == s
