# this will be shared between isolated parts and exposed parts

import os
import tempfile
import pickle
import weakref
import atexit


class TempFileManager:

    _instances = weakref.WeakSet()  # Track instances for cleanup on program exit

    def __init__(self, input_file=None, output_file=None, use_shm=True, file_permissions=0o644, slave = False):
        """
        Initialize the TempFileManager.
        :param input_file: Path to the input file (optional).
        :param output_file: Path to the output file (optional).
        :param use_shm: Whether to attempt using /dev/shm for temporary files.
        """
        self.input_file = input_file
        self.output_file = output_file
        self.use_shm = use_shm and os.path.exists('/dev/shm')
        self.file_permissions = file_permissions
        self.slave = slave

        if not self.slave:
            # Register this instance for cleanup
            TempFileManager._instances.add(self)

    @staticmethod
    def slave_init(input_file, output_file):
        """
        Initialize a TempFileManager instance for use with external filenames.
        :param input_file: Path to the input file.
        :param output_file: Path to the output file.
        :return: A TempFileManager instance with the provided filenames.
        """
        return TempFileManager(input_file=input_file, output_file=output_file, slave=True)

    def __enter__(self):
        """
        Create temporary files when entering the context, if not already provided.
        """
        if not self.input_file:
            base_dir = '/dev/shm' if self.use_shm else None
            self.input_file = tempfile.NamedTemporaryFile(delete=False, dir=base_dir, suffix=".pkl").name
        if not self.output_file:
            base_dir = '/dev/shm' if self.use_shm else None
            self.output_file = tempfile.NamedTemporaryFile(delete=False, dir=base_dir, suffix=".pkl").name
        return self

    def dump_input(self, obj):
        """
        Serialize and write the object to the input file.
        :param obj: The object to serialize and write.
        """
        with open(self.input_file, 'wb') as f:
            pickle.dump(obj, f)
        os.chmod(self.input_file, self.file_permissions)

    def read_input(self):
        """
        Read and deserialize the input file.
        :return: The deserialized object.
        """
        with open(self.input_file, 'rb') as f:
            return pickle.load(f)
        
    def dump_output(self, obj):
        """
        Serialize and write the object to the output file.
        :param obj: The object to serialize and write.
        """
        with open(self.output_file, 'wb') as f:
            pickle.dump(obj, f)
        os.chmod(self.output_file, self.file_permissions)

    def read_output(self):
        """
        Read and deserialize the output file.
        :return: The deserialized object.
        """
        with open(self.output_file, 'rb') as f:
            return pickle.load(f)

    def cleanup(self):
        """
        Cleanup temporary files when exiting the context.
        """
        for file_path in [self.input_file, self.output_file]:
            if file_path and os.path.exists(file_path):
                try:
                    os.remove(file_path)
                except OSError:
                    pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cleanup()


# Ensure all instances are cleaned up on exit
@atexit.register
def cleanup_all_temp_files():
    for instance in TempFileManager._instances:
        instance.cleanup()


import unittest
import os
import json
import stat


class TestTempFileManager(unittest.TestCase):

    def test_file_creation_and_cleanup(self):
        """Test that temporary files are created and cleaned up properly."""
        with TempFileManager() as tfm:
            # Check if input and output files are created
            self.assertTrue(os.path.exists(tfm.input_file))
            self.assertTrue(os.path.exists(tfm.output_file))
            input_path = tfm.input_file
            output_path = tfm.output_file
        
        # After context exit, files should be deleted
        self.assertFalse(os.path.exists(input_path))
        self.assertFalse(os.path.exists(output_path))

    def test_dump_and_read(self):
        """Test dumping an object to input and reading it back."""
        obj = {"key": "value", "number": 42}
        with TempFileManager() as tfm:
            # Dump object to input file
            tfm.dump_input(obj)
            
            # Read it back to ensure correctness
            with open(tfm.input_file, 'rb') as f:
                data = pickle.load(f)
            self.assertEqual(data, obj)

    def test_slave_init(self):
        """Test the slave_init static method for external file names."""
        input_file = "/tmp/test_input.json"
        output_file = "/tmp/test_output.json"

        try:
            # Create test input file
            obj = {"key": "test"}
            with open(input_file, 'wb') as f:
                pickle.dump(obj, f)
            
            # Initialize TempFileManager with slave_init
            tfm = TempFileManager.slave_init(input_file, output_file)

            # Read the input file and verify content
            data = tfm.read_input()
            self.assertEqual(data, obj)
            
            # Write to the output file and verify
            updated_obj = {"key": "updated"}
            tfm.dump_output(updated_obj)
            with open(output_file, 'rb') as f:
                output_data = pickle.load(f)
            self.assertEqual(output_data, updated_obj)
        finally:
            # Cleanup test files
            if os.path.exists(input_file):
                os.unlink(input_file)
            if os.path.exists(output_file):
                os.unlink(output_file)

    def test_with_dev_shm(self):
        """Test that /dev/shm is used if available."""
        if os.path.exists('/dev/shm'):
            with TempFileManager() as tfm:
                self.assertTrue(tfm.input_file.startswith('/dev/shm'))
                self.assertTrue(tfm.output_file.startswith('/dev/shm'))

    def test_without_dev_shm(self):
        """Test fallback to default location if /dev/shm is unavailable."""
        with TempFileManager(use_shm=False) as tfm:
            self.assertFalse(tfm.input_file.startswith('/dev/shm'))
            self.assertFalse(tfm.output_file.startswith('/dev/shm'))

    def test_permissions_after_dump(self):
        """Test that file permissions are correctly set after dumping content."""
        with TempFileManager(file_permissions=0o600) as tfm:
            # Dump content to the input file
            tfm.dump_input({"key": "value"})
            
            # Check that permissions are applied after dumping
            permissions = stat.S_IMODE(os.stat(tfm.input_file).st_mode)
            self.assertEqual(permissions, 0o600)

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