"""Tests for container bracket-access CLI syntax.

Tests the uniform [] bracket syntax for accessing elements in List and Dict/Mapping
container fields from the CLI. Covers:
- Type system fixes (Sequence/Mapping recognition)
- Container path parsing
- End-to-end CLI parsing with brackets
- YAML config file loading for container fields
- Deep nesting: arbitrary mapping + list + dataclass + poly-dataclass combinations
- Error handling
"""

import sys
import tempfile
from dataclasses import field
from enum import Enum
from pathlib import Path
from typing import Dict, List, Mapping, Sequence
from unittest.mock import patch

import pytest

from vidhi import frozen_dataclass, parse_cli_args
from vidhi.base_poly_config import BasePolyConfig
from vidhi.cli_parser import parse_container_path
from vidhi.utils import (
    _build_container_info,
    is_container,
    is_dict,
    is_list,
    is_valid_container_key_type,
    resolve_container_value,
)

# =============================================================================
# Test Config Dataclasses
# =============================================================================


@frozen_dataclass
class GPULocation:
    node_ip: str = "0.0.0.0"
    gpu_id: int = 0


@frozen_dataclass
class ResourceAllocation:
    resource_mapping: List[GPULocation] = None

    def __post_init__(self):
        if self.resource_mapping is None:
            object.__setattr__(self, "resource_mapping", [])


@frozen_dataclass
class SimpleConfig:
    name: str = "default"
    ids: List[int] = None
    port_map: Dict[str, int] = None
    gpus: List[GPULocation] = None

    def __post_init__(self):
        if self.ids is None:
            object.__setattr__(self, "ids", [])
        if self.port_map is None:
            object.__setattr__(self, "port_map", {})
        if self.gpus is None:
            object.__setattr__(self, "gpus", [])


@frozen_dataclass
class MappingConfig:
    resource_mapping: Mapping[int, ResourceAllocation] = None

    def __post_init__(self):
        if self.resource_mapping is None:
            object.__setattr__(self, "resource_mapping", {})


@frozen_dataclass
class NestedContainerConfig:
    configs: Dict[str, List[GPULocation]] = None

    def __post_init__(self):
        if self.configs is None:
            object.__setattr__(self, "configs", {})


# =============================================================================
# Polymorphic + Deep Nesting Config Dataclasses
# =============================================================================


class StorageType(Enum):
    LOCAL = "local"
    S3 = "s3"


@frozen_dataclass
class BaseStorageConfig(BasePolyConfig):
    @classmethod
    def get_type(cls):
        raise NotImplementedError


@frozen_dataclass
class LocalStorageConfig(BaseStorageConfig):
    path: str = "/tmp"

    @classmethod
    def get_type(cls):
        return StorageType.LOCAL


@frozen_dataclass
class S3StorageConfig(BaseStorageConfig):
    bucket: str = "default"
    region: str = "us-east-1"

    @classmethod
    def get_type(cls):
        return StorageType.S3


@frozen_dataclass
class WorkerConfig:
    name: str = "worker"
    gpus: List[GPULocation] = None
    storage: BaseStorageConfig = field(default_factory=LocalStorageConfig)

    def __post_init__(self):
        if self.gpus is None:
            object.__setattr__(self, "gpus", [])


@frozen_dataclass
class ClusterConfig:
    workers: Mapping[int, WorkerConfig] = None
    gpu_pools: Dict[str, List[GPULocation]] = None

    def __post_init__(self):
        if self.workers is None:
            object.__setattr__(self, "workers", {})
        if self.gpu_pools is None:
            object.__setattr__(self, "gpu_pools", {})


@frozen_dataclass
class DeploymentConfig:
    name: str = "prod"
    clusters: List[ClusterConfig] = None
    storage_options: Dict[str, BaseStorageConfig] = None

    def __post_init__(self):
        if self.clusters is None:
            object.__setattr__(self, "clusters", [])
        if self.storage_options is None:
            object.__setattr__(self, "storage_options", {})


# =============================================================================
# Phase 1: Type System Tests
# =============================================================================


class TestTypeSystemFixes:
    """Test that Sequence and Mapping are recognized as list/dict types."""

    def test_is_list_with_list(self):
        assert is_list(List[int])

    def test_is_list_with_sequence(self):
        assert is_list(Sequence[int])

    def test_is_dict_with_dict(self):
        assert is_dict(Dict[str, int])

    def test_is_dict_with_mapping(self):
        assert is_dict(Mapping[str, int])

    def test_is_container_list(self):
        assert is_container(List[int])

    def test_is_container_mapping(self):
        assert is_container(Mapping[int, str])

    def test_is_container_primitive(self):
        assert not is_container(int)
        assert not is_container(str)

    def test_valid_container_key_types(self):
        assert is_valid_container_key_type(str)
        assert is_valid_container_key_type(int)
        assert is_valid_container_key_type(float)
        assert is_valid_container_key_type(bool)

    def test_invalid_container_key_types(self):
        assert not is_valid_container_key_type(list)
        assert not is_valid_container_key_type(dict)
        assert not is_valid_container_key_type(GPULocation)


# =============================================================================
# Phase 2: Container Info Tests
# =============================================================================


class TestBuildContainerInfo:
    """Test ContainerFieldInfo construction from type annotations."""

    def test_list_of_primitive(self):
        info = _build_container_info(List[int])
        assert info.container_kind == "list"
        assert info.key_type is int
        assert info.value_type is int
        assert not info.value_is_dataclass
        assert info.nested_container is None

    def test_list_of_dataclass(self):
        info = _build_container_info(List[GPULocation])
        assert info.container_kind == "list"
        assert info.key_type is int
        assert info.value_type is GPULocation
        assert info.value_is_dataclass
        assert info.nested_container is None

    def test_dict_of_primitive(self):
        info = _build_container_info(Dict[str, int])
        assert info.container_kind == "dict"
        assert info.key_type is str
        assert info.value_type is int
        assert not info.value_is_dataclass
        assert info.nested_container is None

    def test_mapping_of_dataclass(self):
        info = _build_container_info(Mapping[int, ResourceAllocation])
        assert info.container_kind == "dict"
        assert info.key_type is int
        assert info.value_type is ResourceAllocation
        assert info.value_is_dataclass
        assert info.nested_container is None

    def test_nested_dict_of_list(self):
        info = _build_container_info(Dict[str, List[GPULocation]])
        assert info.container_kind == "dict"
        assert info.key_type is str
        assert info.nested_container is not None
        assert info.nested_container.container_kind == "list"
        assert info.nested_container.value_type is GPULocation
        assert info.nested_container.value_is_dataclass

    def test_invalid_key_type_raises(self):
        with pytest.raises(ValueError, match="key type must be a primitive"):
            _build_container_info(Dict[GPULocation, int])


# =============================================================================
# Container Path Parsing Tests
# =============================================================================


class TestParseContainerPath:
    """Test parse_container_path() splits CLI names correctly."""

    def test_no_brackets(self):
        base, segments = parse_container_path("simple_field")
        assert base == "simple_field"
        assert segments == []

    def test_simple_list_access(self):
        base, segments = parse_container_path("gpus[0].node_ip")
        assert base == "gpus"
        assert segments == [("key", "0"), ("field", "node_ip")]

    def test_dict_access(self):
        base, segments = parse_container_path("port_map[http]")
        assert base == "port_map"
        assert segments == [("key", "http")]

    def test_nested_brackets(self):
        base, segments = parse_container_path("configs[train][0].lr")
        assert base == "configs"
        assert segments == [("key", "train"), ("key", "0"), ("field", "lr")]

    def test_dotted_base_with_brackets(self):
        base, segments = parse_container_path("engine.gpus[0].node_ip")
        assert base == "engine.gpus"
        assert segments == [("key", "0"), ("field", "node_ip")]

    def test_bracket_only(self):
        base, segments = parse_container_path("ids[0]")
        assert base == "ids"
        assert segments == [("key", "0")]

    def test_deep_nesting(self):
        base, segments = parse_container_path("rm[0].alloc[1].node_ip")
        assert base == "rm"
        assert segments == [
            ("key", "0"),
            ("field", "alloc"),
            ("key", "1"),
            ("field", "node_ip"),
        ]


# =============================================================================
# Resolve Container Value Tests
# =============================================================================


class TestResolveContainerValue:
    """Test resolve_container_value() converts raw dicts to typed instances."""

    def test_list_of_primitives_from_indexed_dict(self):
        info = _build_container_info(List[int])
        result = resolve_container_value({0: "10", 1: "20", 2: "30"}, info)
        assert result == [10, 20, 30]

    def test_list_of_dataclass_from_indexed_dict(self):
        info = _build_container_info(List[GPULocation])
        result = resolve_container_value(
            {0: {"node_ip": "10.0.0.1", "gpu_id": 0}}, info
        )
        assert len(result) == 1
        assert isinstance(result[0], GPULocation)
        assert result[0].node_ip == "10.0.0.1"
        assert result[0].gpu_id == 0

    def test_dict_of_primitives(self):
        info = _build_container_info(Dict[str, int])
        result = resolve_container_value({"http": "8080", "grpc": "50051"}, info)
        assert result == {"http": 8080, "grpc": 50051}

    def test_mapping_of_dataclass(self):
        info = _build_container_info(Mapping[int, ResourceAllocation])
        result = resolve_container_value(
            {"0": {"resource_mapping": [{"node_ip": "10.0.0.1", "gpu_id": 0}]}},
            info,
        )
        assert 0 in result
        assert isinstance(result[0], ResourceAllocation)
        assert len(result[0].resource_mapping) == 1
        assert result[0].resource_mapping[0].node_ip == "10.0.0.1"

    def test_sparse_indices_raises(self):
        info = _build_container_info(List[int])
        with pytest.raises(ValueError, match="Sparse list indices"):
            resolve_container_value({0: "10", 5: "50"}, info)

    def test_none_returns_none(self):
        info = _build_container_info(List[int])
        assert resolve_container_value(None, info) is None


# =============================================================================
# End-to-End CLI Tests
# =============================================================================


class TestContainerCLI:
    """End-to-end tests for bracket-access CLI syntax."""

    def test_list_primitive_shorthand(self):
        """Existing multi-value shorthand still works."""
        test_args = ["script.py", "--ids", "1", "2", "3"]
        with patch.object(sys, "argv", test_args):
            config = parse_cli_args(SimpleConfig)
        assert config.ids == [1, 2, 3]

    def test_list_primitive_bracket(self):
        """Bracket syntax works for primitive lists."""
        test_args = ["script.py", "--ids[0]", "10", "--ids[1]", "20"]
        with patch.object(sys, "argv", test_args):
            config = parse_cli_args(SimpleConfig)
        assert config.ids == [10, 20]

    def test_dict_primitive_bracket(self):
        """Bracket syntax for Dict[str, int]."""
        test_args = [
            "script.py",
            "--port_map[http]",
            "8080",
            "--port_map[grpc]",
            "50051",
        ]
        with patch.object(sys, "argv", test_args):
            config = parse_cli_args(SimpleConfig)
        assert config.port_map == {"http": 8080, "grpc": 50051}

    def test_list_dataclass_bracket(self):
        """Bracket syntax for List[DataClass]."""
        test_args = [
            "script.py",
            "--gpus[0].node_ip",
            "10.0.0.1",
            "--gpus[0].gpu_id",
            "0",
            "--gpus[1].node_ip",
            "10.0.0.2",
            "--gpus[1].gpu_id",
            "1",
        ]
        with patch.object(sys, "argv", test_args):
            config = parse_cli_args(SimpleConfig)
        assert len(config.gpus) == 2
        assert config.gpus[0].node_ip == "10.0.0.1"
        assert config.gpus[0].gpu_id == 0
        assert config.gpus[1].node_ip == "10.0.0.2"
        assert config.gpus[1].gpu_id == 1

    def test_mapping_bracket(self):
        """Bracket syntax for Mapping[int, DataClass]."""
        test_args = [
            "script.py",
            "--resource_mapping[0].resource_mapping[0].node_ip",
            "10.0.0.1",
            "--resource_mapping[0].resource_mapping[0].gpu_id",
            "0",
        ]
        with patch.object(sys, "argv", test_args):
            config = parse_cli_args(MappingConfig)
        assert 0 in config.resource_mapping
        alloc = config.resource_mapping[0]
        assert isinstance(alloc, ResourceAllocation)
        assert len(alloc.resource_mapping) == 1
        assert alloc.resource_mapping[0].node_ip == "10.0.0.1"
        assert alloc.resource_mapping[0].gpu_id == 0

    def test_bracket_with_equals_syntax(self):
        """Bracket with --flag=value syntax."""
        test_args = ["script.py", "--port_map[http]=8080"]
        with patch.object(sys, "argv", test_args):
            config = parse_cli_args(SimpleConfig)
        assert config.port_map == {"http": 8080}


class TestContainerErrors:
    """Test error handling for container bracket access."""

    def test_bracket_on_non_container(self):
        """Error when using brackets on a scalar field."""
        test_args = ["script.py", "--name[0]", "test"]
        with patch.object(sys, "argv", test_args):
            with pytest.raises(SystemExit):
                parse_cli_args(SimpleConfig)


# =============================================================================
# Deep Nesting Tests (CLI bracket + YAML file)
# =============================================================================


class TestDeepNesting:
    """Test deeply nested container + dataclass + poly-dataclass combinations.

    Each pattern is tested via both CLI bracket syntax and YAML config file.
    """

    # --- Depth 3: Dict[str, List[DataClass]] ---

    def test_dict_str_list_dataclass_cli(self):
        """Dict[str, List[DataClass]] via CLI bracket syntax (3 levels deep)."""
        test_args = [
            "script.py",
            "--configs[pool_a][0].node_ip",
            "10.0.0.1",
            "--configs[pool_a][0].gpu_id",
            "0",
            "--configs[pool_a][1].node_ip",
            "10.0.0.2",
            "--configs[pool_a][1].gpu_id",
            "1",
            "--configs[pool_b][0].node_ip",
            "10.0.0.3",
            "--configs[pool_b][0].gpu_id",
            "2",
        ]
        with patch.object(sys, "argv", test_args):
            config = parse_cli_args(NestedContainerConfig)
        assert len(config.configs["pool_a"]) == 2
        assert config.configs["pool_a"][0].node_ip == "10.0.0.1"
        assert config.configs["pool_a"][0].gpu_id == 0
        assert config.configs["pool_a"][1].node_ip == "10.0.0.2"
        assert config.configs["pool_b"][0].gpu_id == 2

    def test_dict_str_list_dataclass_yaml(self):
        """Dict[str, List[DataClass]] via YAML config file (3 levels deep)."""
        yaml_content = """\
configs:
  pool_a:
    - node_ip: "10.0.0.1"
      gpu_id: 0
    - node_ip: "10.0.0.2"
      gpu_id: 1
  pool_b:
    - node_ip: "10.0.0.3"
      gpu_id: 2
"""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            f.write(yaml_content)
            yaml_path = f.name
        try:
            with patch.object(sys, "argv", ["script.py", "--config", yaml_path]):
                config = parse_cli_args(NestedContainerConfig)
            assert len(config.configs["pool_a"]) == 2
            assert config.configs["pool_a"][0].node_ip == "10.0.0.1"
            assert config.configs["pool_a"][1].gpu_id == 1
            assert config.configs["pool_b"][0].gpu_id == 2
        finally:
            Path(yaml_path).unlink()

    # --- Depth 3: Mapping[int, DC{List[DC]}] ---

    def test_mapping_dc_with_list_dc_cli(self):
        """Mapping[int, DC-containing-List[DC]] via CLI bracket syntax (3 levels)."""
        test_args = [
            "script.py",
            "--resource_mapping[0].resource_mapping[0].node_ip",
            "10.0.0.1",
            "--resource_mapping[0].resource_mapping[0].gpu_id",
            "0",
            "--resource_mapping[0].resource_mapping[1].node_ip",
            "10.0.0.2",
            "--resource_mapping[0].resource_mapping[1].gpu_id",
            "1",
            "--resource_mapping[1].resource_mapping[0].node_ip",
            "10.0.0.3",
            "--resource_mapping[1].resource_mapping[0].gpu_id",
            "2",
        ]
        with patch.object(sys, "argv", test_args):
            config = parse_cli_args(MappingConfig)
        assert len(config.resource_mapping) == 2
        assert len(config.resource_mapping[0].resource_mapping) == 2
        assert config.resource_mapping[0].resource_mapping[0].node_ip == "10.0.0.1"
        assert config.resource_mapping[0].resource_mapping[1].node_ip == "10.0.0.2"
        assert config.resource_mapping[1].resource_mapping[0].gpu_id == 2

    def test_mapping_dc_with_list_dc_yaml(self):
        """Mapping[int, DC-containing-List[DC]] via YAML config file (3 levels)."""
        yaml_content = """\
resource_mapping:
  0:
    resource_mapping:
      - node_ip: "10.0.0.1"
        gpu_id: 0
      - node_ip: "10.0.0.2"
        gpu_id: 1
  1:
    resource_mapping:
      - node_ip: "10.0.0.3"
        gpu_id: 2
"""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            f.write(yaml_content)
            yaml_path = f.name
        try:
            with patch.object(sys, "argv", ["script.py", "--config", yaml_path]):
                config = parse_cli_args(MappingConfig)
            assert len(config.resource_mapping) == 2
            assert config.resource_mapping[0].resource_mapping[0].node_ip == "10.0.0.1"
            assert config.resource_mapping[0].resource_mapping[1].node_ip == "10.0.0.2"
            assert config.resource_mapping[1].resource_mapping[0].gpu_id == 2
        finally:
            Path(yaml_path).unlink()

    # --- Depth 2: Dict[str, PolyDC] ---

    def test_dict_str_poly_dc_cli(self):
        """Dict[str, PolyDC] via CLI bracket syntax — poly dispatch inside container."""
        test_args = [
            "script.py",
            "--storage_options[primary].type",
            "s3",
            "--storage_options[primary].bucket",
            "my-bucket",
            "--storage_options[primary].region",
            "us-west-2",
            "--storage_options[backup].type",
            "local",
            "--storage_options[backup].path",
            "/mnt/backup",
        ]
        with patch.object(sys, "argv", test_args):
            config = parse_cli_args(DeploymentConfig)
        assert isinstance(config.storage_options["primary"], S3StorageConfig)
        assert config.storage_options["primary"].bucket == "my-bucket"
        assert config.storage_options["primary"].region == "us-west-2"
        assert isinstance(config.storage_options["backup"], LocalStorageConfig)
        assert config.storage_options["backup"].path == "/mnt/backup"

    def test_dict_str_poly_dc_yaml(self):
        """Dict[str, PolyDC] via YAML config file — poly dispatch inside container."""
        yaml_content = """\
storage_options:
  primary:
    type: s3
    bucket: my-bucket
    region: us-west-2
  backup:
    type: local
    path: /mnt/backup
"""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            f.write(yaml_content)
            yaml_path = f.name
        try:
            with patch.object(sys, "argv", ["script.py", "--config", yaml_path]):
                config = parse_cli_args(DeploymentConfig)
            assert isinstance(config.storage_options["primary"], S3StorageConfig)
            assert config.storage_options["primary"].bucket == "my-bucket"
            assert config.storage_options["primary"].region == "us-west-2"
            assert isinstance(config.storage_options["backup"], LocalStorageConfig)
            assert config.storage_options["backup"].path == "/mnt/backup"
        finally:
            Path(yaml_path).unlink()

    # --- Depth 4: Mapping[int, DC{List[DC], PolyDC}] ---

    def test_mapping_dc_with_list_and_poly_cli(self):
        """Mapping[int, DC{List[DC], PolyDC}] via CLI bracket syntax (4 levels)."""
        test_args = [
            "script.py",
            "--workers[0].name",
            "w0",
            "--workers[0].gpus[0].node_ip",
            "10.0.0.1",
            "--workers[0].gpus[0].gpu_id",
            "0",
            "--workers[0].storage.type",
            "s3",
            "--workers[0].storage.bucket",
            "w0-bucket",
            "--workers[1].name",
            "w1",
            "--workers[1].gpus[0].node_ip",
            "10.0.0.2",
            "--workers[1].gpus[0].gpu_id",
            "1",
            "--workers[1].storage.type",
            "local",
            "--workers[1].storage.path",
            "/data/w1",
        ]
        with patch.object(sys, "argv", test_args):
            config = parse_cli_args(ClusterConfig)
        assert config.workers[0].name == "w0"
        assert len(config.workers[0].gpus) == 1
        assert config.workers[0].gpus[0].node_ip == "10.0.0.1"
        assert isinstance(config.workers[0].storage, S3StorageConfig)
        assert config.workers[0].storage.bucket == "w0-bucket"
        assert config.workers[1].name == "w1"
        assert isinstance(config.workers[1].storage, LocalStorageConfig)
        assert config.workers[1].storage.path == "/data/w1"

    def test_mapping_dc_with_list_and_poly_yaml(self):
        """Mapping[int, DC{List[DC], PolyDC}] via YAML config file (4 levels)."""
        yaml_content = """\
workers:
  0:
    name: w0
    gpus:
      - node_ip: "10.0.0.1"
        gpu_id: 0
    storage:
      type: s3
      bucket: w0-bucket
  1:
    name: w1
    gpus:
      - node_ip: "10.0.0.2"
        gpu_id: 1
    storage:
      type: local
      path: /data/w1
"""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            f.write(yaml_content)
            yaml_path = f.name
        try:
            with patch.object(sys, "argv", ["script.py", "--config", yaml_path]):
                config = parse_cli_args(ClusterConfig)
            assert config.workers[0].name == "w0"
            assert config.workers[0].gpus[0].node_ip == "10.0.0.1"
            assert isinstance(config.workers[0].storage, S3StorageConfig)
            assert config.workers[0].storage.bucket == "w0-bucket"
            assert config.workers[1].name == "w1"
            assert isinstance(config.workers[1].storage, LocalStorageConfig)
            assert config.workers[1].storage.path == "/data/w1"
        finally:
            Path(yaml_path).unlink()

    # --- Depth 5: List[DC{Mapping[int, DC{List[DC], PolyDC}], Dict[str, List[DC]]}] ---

    def test_list_cluster_full_cli(self):
        """List[ClusterConfig] via CLI bracket syntax (5 levels deep)."""
        test_args = [
            "script.py",
            "--clusters[0].workers[0].name",
            "w0",
            "--clusters[0].workers[0].gpus[0].node_ip",
            "10.0.0.1",
            "--clusters[0].workers[0].gpus[0].gpu_id",
            "0",
            "--clusters[0].gpu_pools[pool_a][0].node_ip",
            "10.0.0.2",
            "--clusters[0].gpu_pools[pool_a][0].gpu_id",
            "1",
        ]
        with patch.object(sys, "argv", test_args):
            config = parse_cli_args(DeploymentConfig)
        assert len(config.clusters) == 1
        cluster = config.clusters[0]
        assert cluster.workers[0].name == "w0"
        assert cluster.workers[0].gpus[0].node_ip == "10.0.0.1"
        assert cluster.gpu_pools["pool_a"][0].node_ip == "10.0.0.2"
        assert cluster.gpu_pools["pool_a"][0].gpu_id == 1

    def test_list_cluster_full_yaml(self):
        """List[ClusterConfig] via YAML config file (5 levels deep)."""
        yaml_content = """\
clusters:
  - workers:
      0:
        name: w0
        gpus:
          - node_ip: "10.0.0.1"
            gpu_id: 0
        storage:
          type: s3
          bucket: cluster-bucket
    gpu_pools:
      pool_a:
        - node_ip: "10.0.0.2"
          gpu_id: 1
"""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            f.write(yaml_content)
            yaml_path = f.name
        try:
            with patch.object(sys, "argv", ["script.py", "--config", yaml_path]):
                config = parse_cli_args(DeploymentConfig)
            assert len(config.clusters) == 1
            cluster = config.clusters[0]
            assert cluster.workers[0].name == "w0"
            assert cluster.workers[0].gpus[0].node_ip == "10.0.0.1"
            assert isinstance(cluster.workers[0].storage, S3StorageConfig)
            assert cluster.workers[0].storage.bucket == "cluster-bucket"
            assert cluster.gpu_pools["pool_a"][0].node_ip == "10.0.0.2"
        finally:
            Path(yaml_path).unlink()

    # --- Comprehensive: Full DeploymentConfig with all field types ---

    def test_full_deployment_yaml(self):
        """Full DeploymentConfig exercising all nested types via YAML."""
        yaml_content = """\
name: staging
clusters:
  - workers:
      0:
        name: gpu-worker-0
        gpus:
          - node_ip: "10.0.0.1"
            gpu_id: 0
          - node_ip: "10.0.0.1"
            gpu_id: 1
        storage:
          type: s3
          bucket: staging-data
          region: us-west-2
      1:
        name: gpu-worker-1
        gpus:
          - node_ip: "10.0.0.2"
            gpu_id: 0
        storage:
          type: local
          path: /mnt/ssd
    gpu_pools:
      training:
        - node_ip: "10.0.0.1"
          gpu_id: 0
        - node_ip: "10.0.0.1"
          gpu_id: 1
      inference:
        - node_ip: "10.0.0.2"
          gpu_id: 0
storage_options:
  primary:
    type: s3
    bucket: main-storage
    region: us-east-1
  cache:
    type: local
    path: /tmp/cache
"""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            f.write(yaml_content)
            yaml_path = f.name
        try:
            with patch.object(sys, "argv", ["script.py", "--config", yaml_path]):
                config = parse_cli_args(DeploymentConfig)
            # Top level
            assert config.name == "staging"
            # Cluster -> workers
            cluster = config.clusters[0]
            assert cluster.workers[0].name == "gpu-worker-0"
            assert len(cluster.workers[0].gpus) == 2
            assert cluster.workers[0].gpus[1].gpu_id == 1
            assert isinstance(cluster.workers[0].storage, S3StorageConfig)
            assert cluster.workers[0].storage.bucket == "staging-data"
            assert cluster.workers[1].name == "gpu-worker-1"
            assert isinstance(cluster.workers[1].storage, LocalStorageConfig)
            assert cluster.workers[1].storage.path == "/mnt/ssd"
            # Cluster -> gpu_pools
            assert len(cluster.gpu_pools["training"]) == 2
            assert len(cluster.gpu_pools["inference"]) == 1
            assert cluster.gpu_pools["training"][0].node_ip == "10.0.0.1"
            # Top-level storage_options
            assert isinstance(config.storage_options["primary"], S3StorageConfig)
            assert config.storage_options["primary"].region == "us-east-1"
            assert isinstance(config.storage_options["cache"], LocalStorageConfig)
            assert config.storage_options["cache"].path == "/tmp/cache"
        finally:
            Path(yaml_path).unlink()
