import concurrent.futures
from typing import Optional
import os

def _worker_init_disable_mathjax():
    """Initializer for worker processes to ensure Plotly does not try to use MathJax.
    Newer Plotly uses pio.defaults.mathjax; keep a safe fallback for older versions.
    """
    try:
        import plotly.io as pio  # local import inside worker
        # Preferred API (Plotly 5.10+)
        if hasattr(pio, "defaults"):
            pio.defaults.mathjax = None
        else:  # Fallback for very old Plotly
            if hasattr(pio, "kaleido") and hasattr(pio.kaleido, "scope"):
                pio.kaleido.scope.mathjax = None
    except Exception:
        # Do not fail worker startup if Plotly import/config fails; rendering will raise later if needed
        pass


class BackgroundProcessManager:
    _instance: "BackgroundProcessManager | None" = None

    def __init__(self):
        if not hasattr(self, 'executor'):
            # Use a process pool to bypass GIL; good for CPU-bound rendering tasks
            cpu_count = os.cpu_count()

            max_workers = max(1, int(cpu_count / 3))

            self.executor = concurrent.futures.ProcessPoolExecutor(
                max_workers=max_workers, initializer=_worker_init_disable_mathjax
            )
            self.futures: list[concurrent.futures.Future] = []
            self.parallel_enabled: bool = True

    @classmethod
    def get_instance(cls):
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    def set_parallel_enabled(self, enabled: bool):
        self.parallel_enabled = enabled

    def submit_task(self, func, *args, **kwargs):
        if not self.parallel_enabled:
            # Run synchronously when parallel is disabled (useful for tests/debug)
            try:
                result = func(*args, **kwargs)
            except Exception as e:
                # Mirror behavior of futures by raising here
                raise e
            return result

        future = self.executor.submit(func, *args, **kwargs)
        self.futures.append(future)
        return future

    def wait_all(self):
        # Wait for all futures to complete
        if not self.futures:
            return

        for future in concurrent.futures.as_completed(self.futures):
            # Retrieve result to raise any exceptions that occurred during execution
            try:
                future.result()
            except Exception as e:
                print(f"Background task failed: {e}")
                # Depending on requirements, we might want to re-raise or log
                raise e
        
        # Clear futures after completion
        self.futures.clear()

    def reset(self):
        """Clear state between report runs (called from cleanup)."""
        self.futures.clear()


# Global instance
manager = BackgroundProcessManager.get_instance()
