import os
import shutil
import subprocess
import unittest


class TestBase(unittest.TestCase):
    def setUp(self):
        # Ensure output directory exists
        os.makedirs("test_output", exist_ok=True)

        # Set up environment variables
        self.env = os.environ.copy()
        self.project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
        self.env["PYTHONPATH"] = f"{self.project_root}:{self.env.get('PYTHONPATH', '')}"

        # List to store files and directories to clean up
        self._outputs_to_clean = []

    def run_script(self, script_name, input_file, output_file):
        """Run the specified script with given input and output files/directories."""
        # Register output for cleanup
        self.register_output(output_file)

        subprocess.run(
            ["uv", "run", "python", f"scripts/{script_name}", input_file, output_file], check=True, env=self.env
        )

    def run_script_with_output(self, script_name, input_file):
        """Run the specified script with given input and capture stdout output."""
        return subprocess.run(
            ["uv", "run", "python", f"scripts/{script_name}", input_file],
            check=True,
            capture_output=True,
            text=True,
            env=self.env,
        )

    def verify_output_file(self, output_file):
        """Verify that the output file exists and is not empty."""
        self.assertTrue(os.path.exists(output_file), f"Output file {output_file} was not created")
        self.assertGreater(os.path.getsize(output_file), 0, f"Output file {output_file} is empty")

    def verify_output_dir(self, output_dir):
        """Verify that the output directory exists and is not empty."""
        self.assertTrue(os.path.exists(output_dir), f"Output directory {output_dir} was not created")
        self.assertGreater(len(os.listdir(output_dir)), 0, f"Output directory {output_dir} is empty")

    def register_output(self, output_path):
        """Register a file or directory for cleanup after the test."""
        if output_path not in self._outputs_to_clean:
            self._outputs_to_clean.append(output_path)

    def register_outputs(self, output_paths):
        """Register multiple files or directories for cleanup after the test."""
        for output_path in output_paths:
            self.register_output(output_path)

    def tearDown(self):
        # Clean up all registered output files and directories
        for output_path in reversed(self._outputs_to_clean):
            if os.path.exists(output_path):
                if os.path.isdir(output_path):
                    shutil.rmtree(output_path)
                else:
                    os.remove(output_path)

        # Also clean up any numbered PNG files that might be generated by some tests
        for file in os.listdir("test_output"):
            if file.startswith("test_") and file.endswith(".png"):
                file_path = os.path.join("test_output", file)
                if os.path.exists(file_path):
                    os.remove(file_path)
