"""Tests for infrastructure spec models (Hetzner workers, AWS Batch, Docker build)."""

from __future__ import annotations

import pytest
from pydantic import ValidationError

from prisme.spec.infrastructure import (
    AWSBatchConfig,
    AWSBatchJobSpec,
    DockerBuildServiceConfig,
    HetznerWorkerJobSpec,
    HetznerWorkersConfig,
)


class TestHetznerWorkerJobSpec:
    """Tests for HetznerWorkerJobSpec model."""

    def test_minimal(self) -> None:
        job = HetznerWorkerJobSpec(name="test", command="python run.py")
        assert job.name == "test"
        assert job.restart_policy == "on-failure"
        assert job.cpu == 2
        assert job.memory_mb == 4096

    def test_all_fields(self) -> None:
        job = HetznerWorkerJobSpec(
            name="full-job",
            command="python -m worker",
            schedule="0 * * * *",
            server_type="cx32",
            location="hel1",
            image="ubuntu-22.04",
            cpu=4,
            memory_mb=8192,
            restart_policy="always",
            max_restarts=5,
            timeout_seconds=7200,
            environment={"KEY": "value"},
            docker_image="registry/img:v1",
        )
        assert job.schedule == "0 * * * *"
        assert job.restart_policy == "always"

    def test_invalid_restart_policy(self) -> None:
        with pytest.raises(ValidationError):
            HetznerWorkerJobSpec(
                name="bad",
                command="cmd",
                restart_policy="invalid",  # type: ignore[arg-type]
            )


class TestHetznerWorkersConfig:
    """Tests for HetznerWorkersConfig model."""

    def test_defaults(self) -> None:
        config = HetznerWorkersConfig()
        assert config.enabled is False
        assert config.default_server_type == "cx22"
        assert config.jobs == []

    def test_with_jobs(self) -> None:
        config = HetznerWorkersConfig(
            enabled=True,
            jobs=[
                HetznerWorkerJobSpec(name="a", command="cmd_a"),
                HetznerWorkerJobSpec(name="b", command="cmd_b"),
            ],
        )
        assert len(config.jobs) == 2


class TestAWSBatchJobSpec:
    """Tests for AWSBatchJobSpec model."""

    def test_minimal(self) -> None:
        job = AWSBatchJobSpec(name="etl", command=["python", "-m", "etl"])
        assert job.vcpus == 1
        assert job.memory_mb == 2048
        assert job.platform == "FARGATE"
        assert job.retry_attempts == 1

    def test_scheduled(self) -> None:
        job = AWSBatchJobSpec(
            name="hourly",
            command=["python", "run.py"],
            schedule="rate(1 hour)",
        )
        assert job.schedule == "rate(1 hour)"

    def test_ec2_platform(self) -> None:
        job = AWSBatchJobSpec(
            name="gpu-job",
            command=["python", "train.py"],
            platform="EC2",
            vcpus=8,
            memory_mb=32768,
        )
        assert job.platform == "EC2"
        assert job.vcpus == 8

    def test_invalid_platform(self) -> None:
        with pytest.raises(ValidationError):
            AWSBatchJobSpec(
                name="bad",
                command=["cmd"],
                platform="LAMBDA",  # type: ignore[arg-type]
            )


class TestAWSBatchConfig:
    """Tests for AWSBatchConfig model."""

    def test_defaults(self) -> None:
        config = AWSBatchConfig()
        assert config.enabled is False
        assert config.region == "eu-west-1"
        assert config.max_vcpus == 16
        assert config.jobs == []

    def test_with_jobs(self) -> None:
        config = AWSBatchConfig(
            enabled=True,
            region="us-east-1",
            jobs=[
                AWSBatchJobSpec(name="job1", command=["cmd1"]),
            ],
        )
        assert config.region == "us-east-1"
        assert len(config.jobs) == 1


class TestDockerBuildServiceConfig:
    """Tests for DockerBuildServiceConfig model."""

    def test_defaults(self) -> None:
        config = DockerBuildServiceConfig()
        assert config.provider == "github"
        assert config.registry == "ghcr.io"
        assert config.platforms == ["linux/amd64"]

    def test_hetzner_provider(self) -> None:
        config = DockerBuildServiceConfig(
            provider="hetzner",
            build_server_type="cx42",
            build_location="nbg1",
        )
        assert config.provider == "hetzner"
        assert config.build_server_type == "cx42"

    def test_aws_provider(self) -> None:
        config = DockerBuildServiceConfig(
            provider="aws",
            aws_region="us-east-1",
            ecr_repository="my-app",
            registry="123456.dkr.ecr.us-east-1.amazonaws.com",
        )
        assert config.provider == "aws"
        assert config.ecr_repository == "my-app"

    def test_invalid_provider(self) -> None:
        with pytest.raises(ValidationError):
            DockerBuildServiceConfig(provider="azure")  # type: ignore[arg-type]

    def test_multi_platform(self) -> None:
        config = DockerBuildServiceConfig(
            platforms=["linux/amd64", "linux/arm64"],
        )
        assert len(config.platforms) == 2

    def test_build_args(self) -> None:
        config = DockerBuildServiceConfig(
            build_args={"NODE_ENV": "production", "PY_VER": "3.13"},
        )
        assert config.build_args["NODE_ENV"] == "production"
