# ____________________________________________________________________________________
#
# Pyomo: Python Optimization Modeling Objects
# Copyright (c) 2008-2026 National Technology and Engineering Solutions of Sandia, LLC
# Under the terms of Contract DE-NA0003525 with National Technology and Engineering
# Solutions of Sandia, LLC, the U.S. Government retains certain rights in this
# software.  This software is distributed under the 3-clause BSD License.
# ____________________________________________________________________________________

import os
import pyomo.common.unittest as unittest

from pyomo.common.env import CtypesEnviron


class TestCtypesEnviron(unittest.TestCase):
    def test_temp_env_str(self):
        orig_env = CtypesEnviron()
        orig_env_has_1 = 'TEST_ENV_1' in orig_env
        orig_env_has_2 = 'TEST_ENV_2' in orig_env

        # Ensure that TEST_ENV_1 is not in the environment
        if 'TEST_ENV_1' in orig_env:
            del orig_env['TEST_ENV_1']
        orig_env['TEST_ENV_2'] = "test value: 2"

        self.assertIsNone(os.environ.get('TEST_ENV_1', None))
        self.assertEqual(os.environ['TEST_ENV_2'], "test value: 2")
        for interface in orig_env.interfaces:
            self.assertIsNone(interface.dll.wgetenv(u'TEST_ENV_1'))
            self.assertIsNone(interface.dll.getenv(b'TEST_ENV_1'))
            self.assertEqual(interface.dll.wgetenv(u'TEST_ENV_2'), u"test value: 2")
            self.assertEqual(interface.dll.getenv(b'TEST_ENV_2'), b"test value: 2")

        with CtypesEnviron(TEST_ENV_1="test value: 1") as env:
            self.assertEqual(os.environ['TEST_ENV_1'], "test value: 1")
            self.assertEqual(os.environ['TEST_ENV_2'], "test value: 2")
            for interface in env.interfaces:
                self.assertEqual(interface.dll.wgetenv(u'TEST_ENV_1'), u"test value: 1")
                self.assertEqual(interface.dll.getenv(b'TEST_ENV_1'), b"test value: 1")
                self.assertEqual(interface.dll.wgetenv(u'TEST_ENV_2'), u"test value: 2")
                self.assertEqual(interface.dll.getenv(b'TEST_ENV_2'), b"test value: 2")

            del env['TEST_ENV_2']
            self.assertIsNone(os.environ.get('TEST_ENV_2', None))
            for interface in env.interfaces:
                self.assertIsNone(interface.dll.wgetenv(u'TEST_ENV_2'))
                self.assertIsNone(interface.dll.getenv(b'TEST_ENV_2'))

        self.assertIsNone(os.environ.get('TEST_ENV_1', None))
        self.assertEqual(os.environ['TEST_ENV_2'], "test value: 2")
        for interface in orig_env.interfaces:
            self.assertIsNone(interface.dll.wgetenv(u'TEST_ENV_1'))
            self.assertIsNone(interface.dll.getenv(b'TEST_ENV_1'))
            self.assertEqual(interface.dll.wgetenv(u'TEST_ENV_2'), u"test value: 2")
            self.assertEqual(interface.dll.getenv(b'TEST_ENV_2'), b"test value: 2")

        orig_env.restore()
        self.assertEqual(orig_env_has_1, 'TEST_ENV_1' in os.environ)
        self.assertEqual(orig_env_has_2, 'TEST_ENV_2' in os.environ)

    def test_temp_env_unicode(self):
        orig_env = CtypesEnviron()
        orig_env_has_1 = u'TEST_ENV_1' in orig_env
        orig_env_has_2 = u'TEST_ENV_2' in orig_env

        # Ensure that TEST_ENV_1 is not in the environment
        if u'TEST_ENV_1' in orig_env:
            del orig_env[u'TEST_ENV_1']
        orig_env[u'TEST_ENV_2'] = u"test value: 2"

        self.assertIsNone(os.environ.get(u'TEST_ENV_1', None))
        self.assertEqual(os.environ[u'TEST_ENV_2'], "test value: 2")
        for interface in orig_env.interfaces:
            self.assertIsNone(interface.dll.wgetenv(u'TEST_ENV_1'))
            self.assertIsNone(interface.dll.getenv(b'TEST_ENV_1'))
            self.assertEqual(interface.dll.wgetenv(u'TEST_ENV_2'), u"test value: 2")
            self.assertEqual(interface.dll.getenv(b'TEST_ENV_2'), b"test value: 2")

        with CtypesEnviron(TEST_ENV_1=u"test value: 1") as env:
            self.assertEqual(os.environ[u'TEST_ENV_1'], u"test value: 1")
            self.assertEqual(os.environ[u'TEST_ENV_2'], u"test value: 2")
            for interface in env.interfaces:
                self.assertEqual(interface.dll.wgetenv(u'TEST_ENV_1'), u"test value: 1")
                self.assertEqual(interface.dll.getenv(b'TEST_ENV_1'), b"test value: 1")
                self.assertEqual(interface.dll.wgetenv(u'TEST_ENV_2'), u"test value: 2")
                self.assertEqual(interface.dll.getenv(b'TEST_ENV_2'), b"test value: 2")

            del env[u'TEST_ENV_2']
            self.assertIsNone(os.environ.get(u'TEST_ENV_2', None))
            for interface in env.interfaces:
                self.assertIsNone(interface.dll.wgetenv(u'TEST_ENV_2'))
                self.assertIsNone(interface.dll.getenv(b'TEST_ENV_2'))

        self.assertIsNone(os.environ.get(u'TEST_ENV_1', None))
        self.assertEqual(os.environ[u'TEST_ENV_2'], u"test value: 2")
        for interface in orig_env.interfaces:
            self.assertIsNone(interface.dll.wgetenv(u'TEST_ENV_1'))
            self.assertIsNone(interface.dll.getenv(b'TEST_ENV_1'))
            self.assertEqual(interface.dll.wgetenv(u'TEST_ENV_2'), u"test value: 2")
            self.assertEqual(interface.dll.getenv(b'TEST_ENV_2'), b"test value: 2")

        orig_env.restore()
        self.assertEqual(orig_env_has_1, u'TEST_ENV_1' in os.environ)
        self.assertEqual(orig_env_has_2, u'TEST_ENV_2' in os.environ)
