import unittest
import os
from dataclasses import dataclass, field
from unittest.mock import patch
import copy

def str_to_bool(value):
    if value.lower() in ("y", "yes", "t", "true", "on", "1"):
        return 1

class KwargsHandler:
    def to_dict(self):
        return copy.deepcopy(self.__dict__)

    def to_kwargs(self):
        default_dict = self.__class__().to_dict()
        this_dict = self.to_dict()
        return {k: v for k, v in this_dict.items() if default_dict[k] != v}

@dataclass
class TorchDynamoPluginCurrent(KwargsHandler):
    dynamic: bool = field(default=None, metadata={"help": "Whether to use dynamic shape"})

    def __post_init__(self):
        prefix = "ACCELERATE_DYNAMO_"
        if self.dynamic is None:
            self.dynamic = str_to_bool(os.environ.get(prefix + "USE_DYNAMIC", "False")) == 1

    def to_dict(self):
        dynamo_config = copy.deepcopy(self.__dict__)
        return dynamo_config

    def to_kwargs(self):
        kwargs = super().to_kwargs()
        return kwargs

@dataclass
class TorchDynamoPluginModified(KwargsHandler):
    dynamic: bool = field(default=None, metadata={"help": "Whether to use dynamic shape"})

    def __post_init__(self):
        prefix = "ACCELERATE_DYNAMO_"
        if self.dynamic is None:
            env_value = os.environ.get(prefix + "USE_DYNAMIC")
            self.dynamic = (str_to_bool(env_value) == 1) if env_value is not None else None

    def to_dict(self):
        dynamo_config = copy.deepcopy(self.__dict__)
        return dynamo_config

    def to_kwargs(self):
        kwargs = super().to_kwargs()
        return kwargs

class TestTorchDynamoPlugin(unittest.TestCase):
    def test_dynamic_current_implementation(self):
        """Test the current implementation of TorchDynamoPlugin.__post_init__"""
        # Test case 1: Environment variable set to True
        with patch.dict(os.environ, {"ACCELERATE_DYNAMO_USE_DYNAMIC": "True"}):
            plugin = TorchDynamoPluginCurrent()
            self.assertTrue(plugin.dynamic, "Expected dynamic to be True when env is 'True'")

        # Test case 2: Environment variable set to False
        with patch.dict(os.environ, {"ACCELERATE_DYNAMO_USE_DYNAMIC": "False"}):
            plugin = TorchDynamoPluginCurrent()
            self.assertFalse(plugin.dynamic, "Expected dynamic to be False when env is 'False'")

        # Test case 3: Environment variable unset (should be None, but current fails)
        with patch.dict(os.environ, {}, clear=True):
            plugin = TorchDynamoPluginCurrent()
            self.assertFalse(plugin.dynamic, "Current implementation incorrectly sets dynamic to False when env is unset")

    def test_dynamic_modified_implementation(self):
        """Test the modified implementation of TorchDynamoPlugin.__post_init__"""
        # Test case 1: Environment variable set to True
        with patch.dict(os.environ, {"ACCELERATE_DYNAMO_USE_DYNAMIC": "True"}):
            plugin = TorchDynamoPluginModified()
            self.assertTrue(plugin.dynamic, "Expected dynamic to be True when env is 'True'")

        # Test case 2: Environment variable set to False
        with patch.dict(os.environ, {"ACCELERATE_DYNAMO_USE_DYNAMIC": "False"}):
            plugin = TorchDynamoPluginModified()
            self.assertFalse(plugin.dynamic, "Expected dynamic to be False when env is 'False'")

        # Test case 3: Environment variable unset (should be None)
        with patch.dict(os.environ, {}, clear=True):
            plugin = TorchDynamoPluginModified()
            self.assertIsNone(plugin.dynamic, "Expected dynamic to be None when env is unset")

if __name__ == "__main__":
    unittest.main()