"""Tests for cross-variant field leakage, concrete poly type matching, and
merge-with-provided None preservation.

These tests reproduce three bugs that were discovered when polymorphic configs
with variant-specific fields are used with ``with_cli_overrides`` and
``parse_cli_args``.
"""

import sys
import tempfile
from dataclasses import field
from enum import Enum
from pathlib import Path
from typing import Optional
from unittest.mock import patch

from vidhi import (
    BasePolyConfig,
    frozen_dataclass,
    parse_cli_args,
    with_cli_overrides,
)
from vidhi.utils import create_class_from_dict, dataclass_to_dict

# =============================================================================
# Shared fixtures: a realistic multi-variant config hierarchy
# =============================================================================


class ControllerType(Enum):
    LLM = "llm"
    TTS = "tts"


class AllocatorType(Enum):
    LRU = "lru"
    FIFO = "fifo"


@frozen_dataclass
class AbstractModelDeployment(BasePolyConfig):
    """Base model deployment config."""

    model: str = field(default="default-model")
    load_format: str = field(default="safetensors")

    @classmethod
    def get_type(cls):
        raise NotImplementedError()


@frozen_dataclass
class LlmModelDeployment(AbstractModelDeployment):
    """LLM-specific model deployment."""

    seed: int = field(default=0)
    trust_remote_code: bool = field(default=True)
    pipeline_parallel_size: int = field(
        default=1,
        metadata={"aliases": ["pp"]},
    )

    @classmethod
    def get_type(cls) -> ControllerType:
        return ControllerType.LLM


@frozen_dataclass
class TtsModelDeployment(AbstractModelDeployment):
    """TTS-specific model deployment — has a 'device' field that LLM lacks."""

    device: str = field(default="cuda:0")

    @classmethod
    def get_type(cls) -> ControllerType:
        return ControllerType.TTS


@frozen_dataclass
class AbstractAllocator(BasePolyConfig):
    """Base allocator config."""

    @classmethod
    def get_type(cls):
        raise NotImplementedError()


@frozen_dataclass
class LruAllocator(AbstractAllocator):
    max_size: int = field(default=1000)

    @classmethod
    def get_type(cls) -> AllocatorType:
        return AllocatorType.LRU


@frozen_dataclass
class FifoAllocator(AbstractAllocator):
    capacity: int = field(default=500)

    @classmethod
    def get_type(cls) -> AllocatorType:
        return AllocatorType.FIFO


@frozen_dataclass
class AbstractReplicaController(BasePolyConfig):
    """Base replica controller config."""

    model_deployment: AbstractModelDeployment = field(
        default_factory=LlmModelDeployment
    )
    allocator: AbstractAllocator = field(default_factory=LruAllocator)

    @classmethod
    def get_type(cls):
        raise NotImplementedError()


@frozen_dataclass
class LlmReplicaController(AbstractReplicaController):
    max_batch_size: int = field(default=8)

    @classmethod
    def get_type(cls) -> ControllerType:
        return ControllerType.LLM


@frozen_dataclass
class TtsReplicaController(AbstractReplicaController):
    chunk_size: int = field(default=256)

    @classmethod
    def get_type(cls) -> ControllerType:
        return ControllerType.TTS


@frozen_dataclass
class ResourceMapping:
    """Represents allocated GPU resources — may be None when auto-detected."""

    mapping: Optional[dict] = field(default=None)


@frozen_dataclass
class TopLevelConfig:
    """Root config with polymorphic controller and optional resource mapping."""

    controller: AbstractReplicaController = field(default_factory=LlmReplicaController)
    resource_mapping: Optional[ResourceMapping] = field(default=None)
    name: str = field(default="engine")


# =============================================================================
# Bug 1: Cross-variant field leakage
# =============================================================================


class TestCrossVariantFieldLeakage:
    """When the default variant is selected, fields from non-default variants
    (e.g., TTS's ``device``) must NOT leak into the default variant's config.

    Previously, ``_filter_non_selected_variant_defaults`` only filtered fields
    when a *non-default* variant was explicitly selected, leaving non-default
    variant fields in the parsed args when the default was active.
    """

    def test_parse_cli_default_variant_no_leakage(self):
        """parse_cli_args with default variant should not include TTS fields."""
        test_args = ["script.py", "--controller.model_deployment.seed", "42"]
        with patch.object(sys, "argv", test_args):
            config = parse_cli_args(TopLevelConfig)

        assert isinstance(config.controller, LlmReplicaController)
        assert config.controller.model_deployment.seed == 42
        # TTS's 'device' field should not appear
        assert not hasattr(config.controller.model_deployment, "device")

    def test_with_cli_overrides_default_variant_no_leakage(self):
        """with_cli_overrides with default variant should not crash on TTS fields."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(model="my-llm", seed=1)
            )
        )
        test_args = [
            "script.py",
            "--controller.model_deployment.seed",
            "99",
        ]
        with patch.object(sys, "argv", test_args):
            result = with_cli_overrides(base)

        assert isinstance(result.controller, LlmReplicaController)
        assert isinstance(result.controller.model_deployment, LlmModelDeployment)
        assert result.controller.model_deployment.seed == 99
        assert result.controller.model_deployment.model == "my-llm"

    def test_parse_cli_non_default_variant_works(self):
        """Explicitly selecting TTS variant should work and include its fields."""
        test_args = [
            "script.py",
            "--controller.type",
            "tts",
            "--controller.model_deployment.type",
            "tts",
            "--controller.model_deployment.device",
            "cuda:1",
        ]
        with patch.object(sys, "argv", test_args):
            config = parse_cli_args(TopLevelConfig)

        assert isinstance(config.controller, TtsReplicaController)
        assert isinstance(config.controller.model_deployment, TtsModelDeployment)
        assert config.controller.model_deployment.device == "cuda:1"


# =============================================================================
# Bug 2: _match_subclass_by_type doesn't check parent class itself
# =============================================================================


class TestConcretePolyTypeMatching:
    """When ``create_class_from_dict`` encounters a concrete polymorphic class
    used directly as a field type (already resolved at a higher level), the dict
    may contain a ``type`` key. ``_match_subclass_by_type`` must check the
    parent class itself, not just its subclasses.

    Previously, searching only subclasses of a leaf poly class returned no
    candidates, causing a ValueError.
    """

    def test_create_from_dict_concrete_poly_with_type_key(self):
        """create_class_from_dict should handle a concrete poly class with type key."""
        d = {
            "type": "llm",
            "max_batch_size": 16,
            "model_deployment": {
                "type": "llm",
                "model": "test-model",
                "load_format": "safetensors",
                "seed": 42,
                "trust_remote_code": False,
                "pipeline_parallel_size": 2,
            },
            "allocator": {
                "type": "lru",
                "max_size": 2000,
            },
        }
        result = create_class_from_dict(LlmReplicaController, d)
        assert isinstance(result, LlmReplicaController)
        assert result.max_batch_size == 16
        assert result.model_deployment.seed == 42

    def test_roundtrip_concrete_poly(self):
        """Serializing and deserializing a concrete poly config should roundtrip."""
        original = LlmReplicaController(
            model_deployment=LlmModelDeployment(model="roundtrip-model", seed=7),
            allocator=LruAllocator(max_size=500),
            max_batch_size=4,
        )
        d = dataclass_to_dict(original)
        assert "type" in d  # dataclass_to_dict adds type key

        restored = create_class_from_dict(LlmReplicaController, d)
        assert isinstance(restored, LlmReplicaController)
        assert restored.max_batch_size == 4
        assert restored.model_deployment.model == "roundtrip-model"
        assert restored.model_deployment.seed == 7

    def test_with_cli_overrides_roundtrip_with_poly_nesting(self):
        """Full with_cli_overrides roundtrip through nested poly hierarchy."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(model="deep-model"),
                max_batch_size=32,
            )
        )
        test_args = ["script.py", "--controller.max_batch_size", "64"]
        with patch.object(sys, "argv", test_args):
            result = with_cli_overrides(base)

        assert isinstance(result.controller, LlmReplicaController)
        assert result.controller.max_batch_size == 64
        assert result.controller.model_deployment.model == "deep-model"


# =============================================================================
# Bug 3: _merge_with_provided promotes CLI defaults over intentional None
# =============================================================================


class TestMergePreservesNone:
    """``_merge_with_provided`` must preserve intentional ``None`` values from
    the base config. Previously, it used
    ``base_val if base_val is not None else cli_val``, which replaced
    ``None`` with CLI-generated defaults (e.g., an empty dataclass instance).
    """

    def test_optional_none_preserved_after_override(self):
        """Optional field set to None in base should stay None after CLI override."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(model="keep-none"),
            ),
            resource_mapping=None,  # Intentionally None
            name="my-engine",
        )
        test_args = ["script.py", "--name", "overridden"]
        with patch.object(sys, "argv", test_args):
            result = with_cli_overrides(base)

        assert result.name == "overridden"
        # resource_mapping must remain None, not become ResourceMapping(mapping=None)
        assert result.resource_mapping is None

    def test_optional_none_preserved_with_nested_override(self):
        """Overriding a deeply nested field should not promote optional None fields."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(model="nested-test", seed=5),
            ),
            resource_mapping=None,
        )
        test_args = ["script.py", "--controller.model_deployment.seed", "99"]
        with patch.object(sys, "argv", test_args):
            result = with_cli_overrides(base)

        assert result.controller.model_deployment.seed == 99
        assert result.controller.model_deployment.model == "nested-test"
        assert result.resource_mapping is None

    def test_optional_set_value_preserved(self):
        """When Optional field has a value, it should be preserved."""
        base = TopLevelConfig(
            controller=LlmReplicaController(),
            resource_mapping=ResourceMapping(mapping={"gpu0": "cuda:0"}),
        )
        test_args = ["script.py", "--name", "with-mapping"]
        with patch.object(sys, "argv", test_args):
            result = with_cli_overrides(base)

        assert result.name == "with-mapping"
        assert result.resource_mapping is not None
        assert result.resource_mapping.mapping == {"gpu0": "cuda:0"}


# =============================================================================
# Integration: all three bugs together
# =============================================================================


class TestIntegration:
    """End-to-end tests combining all three bug scenarios."""

    def test_full_workflow_cli_override_with_poly_and_optional(self):
        """Simulate the real Vajra workflow: create config, apply CLI overrides."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(
                    model="Qwen/Qwen3-4B",
                    pipeline_parallel_size=1,
                ),
                max_batch_size=8,
            ),
            resource_mapping=None,
            name="vajra-engine",
        )

        # User overrides pipeline_parallel_size via CLI
        test_args = [
            "script.py",
            "--controller.model_deployment.pipeline_parallel_size",
            "2",
        ]
        with patch.object(sys, "argv", test_args):
            result = with_cli_overrides(base)

        # Bug 1: No TTS field leakage
        assert isinstance(result.controller.model_deployment, LlmModelDeployment)
        assert not hasattr(result.controller.model_deployment, "device")

        # Bug 2: Concrete poly roundtrip works
        assert isinstance(result.controller, LlmReplicaController)
        assert result.controller.model_deployment.pipeline_parallel_size == 2

        # Bug 3: Optional None preserved
        assert result.resource_mapping is None

        # Non-overridden values preserved
        assert result.controller.model_deployment.model == "Qwen/Qwen3-4B"
        assert result.controller.max_batch_size == 8
        assert result.name == "vajra-engine"

    def test_full_workflow_yaml_override_with_poly(self):
        """Simulate YAML config override with polymorphic configs."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(model="base-model", seed=1),
                max_batch_size=8,
            ),
            resource_mapping=None,
        )

        yaml_content = """
controller:
    max_batch_size: 16
    model_deployment:
        seed: 42
"""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            f.write(yaml_content)
            yaml_path = f.name

        try:
            test_args = ["script.py", "--config", yaml_path]
            with patch.object(sys, "argv", test_args):
                result = with_cli_overrides(base)

            assert result.controller.max_batch_size == 16
            assert result.controller.model_deployment.seed == 42
            assert result.controller.model_deployment.model == "base-model"
            assert result.resource_mapping is None
        finally:
            Path(yaml_path).unlink()


# =============================================================================
# YAML-only and YAML+CLI override tests
# =============================================================================


def _write_yaml(content: str) -> str:
    """Write YAML content to a temp file and return the path."""
    f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
    f.write(content)
    f.close()
    return f.name


class TestYamlOnlyOverrides:
    """Test all three bug scenarios when overrides come from YAML only."""

    def test_yaml_no_cross_variant_leakage(self):
        """YAML overriding default variant fields must not leak TTS fields."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(model="base-llm", seed=0),
            ),
            resource_mapping=None,
        )
        yaml_path = _write_yaml("""
controller:
    model_deployment:
        seed: 77
        model: yaml-llm
""")
        try:
            with patch.object(sys, "argv", ["s", "--config", yaml_path]):
                result = with_cli_overrides(base)

            assert isinstance(result.controller.model_deployment, LlmModelDeployment)
            assert result.controller.model_deployment.seed == 77
            assert result.controller.model_deployment.model == "yaml-llm"
            assert not hasattr(result.controller.model_deployment, "device")
        finally:
            Path(yaml_path).unlink()

    def test_yaml_select_non_default_variant(self):
        """YAML can select a non-default variant via the type key."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(model="base-llm"),
            ),
        )
        yaml_path = _write_yaml("""
controller:
    type: tts
    model_deployment:
        type: tts
        device: cuda:3
        model: tts-model
""")
        try:
            with patch.object(sys, "argv", ["s", "--config", yaml_path]):
                result = with_cli_overrides(base)

            assert isinstance(result.controller, TtsReplicaController)
            assert isinstance(result.controller.model_deployment, TtsModelDeployment)
            assert result.controller.model_deployment.device == "cuda:3"
            assert result.controller.model_deployment.model == "tts-model"
        finally:
            Path(yaml_path).unlink()

    def test_yaml_preserves_optional_none(self):
        """YAML override of other fields must not promote Optional[None]."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(model="none-test"),
                max_batch_size=4,
            ),
            resource_mapping=None,
            name="original",
        )
        yaml_path = _write_yaml("""
name: from-yaml
controller:
    max_batch_size: 32
""")
        try:
            with patch.object(sys, "argv", ["s", "--config", yaml_path]):
                result = with_cli_overrides(base)

            assert result.name == "from-yaml"
            assert result.controller.max_batch_size == 32
            assert result.resource_mapping is None
        finally:
            Path(yaml_path).unlink()

    def test_yaml_concrete_poly_roundtrip(self):
        """YAML override through nested poly hierarchy must roundtrip cleanly."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(model="roundtrip", seed=1),
                allocator=LruAllocator(max_size=100),
                max_batch_size=2,
            ),
        )
        yaml_path = _write_yaml("""
controller:
    max_batch_size: 64
    allocator:
        max_size: 5000
""")
        try:
            with patch.object(sys, "argv", ["s", "--config", yaml_path]):
                result = with_cli_overrides(base)

            assert isinstance(result.controller, LlmReplicaController)
            assert result.controller.max_batch_size == 64
            assert isinstance(result.controller.allocator, LruAllocator)
            assert result.controller.allocator.max_size == 5000
            assert result.controller.model_deployment.model == "roundtrip"
        finally:
            Path(yaml_path).unlink()


class TestYamlPlusCliOverrides:
    """Test all three bug scenarios with combined YAML + CLI overrides.

    Priority: CLI > YAML > code defaults.
    """

    def test_cli_overrides_yaml_no_leakage(self):
        """CLI overriding a YAML-provided field, default variant, no TTS leakage."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(model="code-model", seed=0),
                max_batch_size=4,
            ),
            resource_mapping=None,
        )
        yaml_path = _write_yaml("""
controller:
    max_batch_size: 16
    model_deployment:
        seed: 10
""")
        try:
            # CLI overrides seed again (CLI > YAML)
            with patch.object(
                sys,
                "argv",
                [
                    "s",
                    "--config",
                    yaml_path,
                    "--controller.model_deployment.seed",
                    "999",
                ],
            ):
                result = with_cli_overrides(base)

            # CLI wins for seed
            assert result.controller.model_deployment.seed == 999
            # YAML wins for max_batch_size
            assert result.controller.max_batch_size == 16
            # Code default wins for model (not in YAML or CLI)
            assert result.controller.model_deployment.model == "code-model"
            # No TTS leakage
            assert isinstance(result.controller.model_deployment, LlmModelDeployment)
            assert not hasattr(result.controller.model_deployment, "device")
            # Optional None preserved
            assert result.resource_mapping is None
        finally:
            Path(yaml_path).unlink()

    def test_yaml_sets_variant_cli_overrides_variant_field(self):
        """YAML selects TTS variant, CLI overrides a TTS-specific field."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(model="base-llm"),
            ),
        )
        yaml_path = _write_yaml("""
controller:
    type: tts
    model_deployment:
        type: tts
        model: yaml-tts
        device: cuda:0
""")
        try:
            with patch.object(
                sys,
                "argv",
                [
                    "s",
                    "--config",
                    yaml_path,
                    "--controller.model_deployment.device",
                    "cuda:7",
                ],
            ):
                result = with_cli_overrides(base)

            assert isinstance(result.controller, TtsReplicaController)
            assert isinstance(result.controller.model_deployment, TtsModelDeployment)
            # CLI wins for device
            assert result.controller.model_deployment.device == "cuda:7"
            # YAML wins for model
            assert result.controller.model_deployment.model == "yaml-tts"
        finally:
            Path(yaml_path).unlink()

    def test_yaml_and_cli_both_override_nested_poly(self):
        """YAML and CLI both touch nested poly fields; check priority and roundtrip."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(
                    model="code-default",
                    seed=1,
                    pipeline_parallel_size=1,
                ),
                allocator=LruAllocator(max_size=100),
                max_batch_size=8,
            ),
            resource_mapping=None,
            name="code-name",
        )
        yaml_path = _write_yaml("""
name: yaml-name
controller:
    max_batch_size: 32
    model_deployment:
        model: yaml-model
        seed: 50
    allocator:
        max_size: 2000
""")
        try:
            with patch.object(
                sys,
                "argv",
                [
                    "s",
                    "--config",
                    yaml_path,
                    "--controller.model_deployment.pipeline_parallel_size",
                    "4",
                    "--name",
                    "cli-name",
                ],
            ):
                result = with_cli_overrides(base)

            # CLI wins
            assert result.name == "cli-name"
            assert result.controller.model_deployment.pipeline_parallel_size == 4
            # YAML wins
            assert result.controller.max_batch_size == 32
            assert result.controller.model_deployment.model == "yaml-model"
            assert result.controller.model_deployment.seed == 50
            assert result.controller.allocator.max_size == 2000
            # Code default wins (not in YAML or CLI)
            assert result.controller.model_deployment.trust_remote_code is True
            # Poly types correct
            assert isinstance(result.controller, LlmReplicaController)
            assert isinstance(result.controller.model_deployment, LlmModelDeployment)
            assert isinstance(result.controller.allocator, LruAllocator)
            # Optional None preserved
            assert result.resource_mapping is None
        finally:
            Path(yaml_path).unlink()

    def test_yaml_partial_cli_fills_rest_preserves_none(self):
        """YAML provides some fields, CLI provides others, None stays None."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                model_deployment=LlmModelDeployment(model="original"),
                max_batch_size=1,
            ),
            resource_mapping=None,
            name="base",
        )
        yaml_path = _write_yaml("""
controller:
    model_deployment:
        model: from-yaml
""")
        try:
            with patch.object(
                sys,
                "argv",
                [
                    "s",
                    "--config",
                    yaml_path,
                    "--controller.max_batch_size",
                    "128",
                ],
            ):
                result = with_cli_overrides(base)

            assert result.controller.model_deployment.model == "from-yaml"
            assert result.controller.max_batch_size == 128
            assert result.name == "base"
            assert result.resource_mapping is None
        finally:
            Path(yaml_path).unlink()

    def test_cli_overrides_yaml_allocator_variant(self):
        """YAML sets one allocator variant, CLI switches to another."""
        base = TopLevelConfig(
            controller=LlmReplicaController(
                allocator=LruAllocator(max_size=100),
            ),
        )
        yaml_path = _write_yaml("""
controller:
    allocator:
        type: lru
        max_size: 500
""")
        try:
            with patch.object(
                sys,
                "argv",
                [
                    "s",
                    "--config",
                    yaml_path,
                    "--controller.allocator.type",
                    "fifo",
                    "--controller.allocator.capacity",
                    "999",
                ],
            ):
                result = with_cli_overrides(base)

            # CLI switches variant to FIFO
            assert isinstance(result.controller.allocator, FifoAllocator)
            assert result.controller.allocator.capacity == 999
        finally:
            Path(yaml_path).unlink()
