import os, sys
from typing import Optional
from PIL import ImageGrab, ImageFilter, Image
from screeninfo import get_monitors
from ws_package.log.writer import LogWriter
from ws_package.utils.screenshot_naming import make_screenshot_filename
from enum import Enum
from pathlib import Path

logger = LogWriter.file(__name__)

# Import mss for Linux (better X11 connection management)
if sys.platform != "darwin":
    try:
        from mss import mss
        MSS_AVAILABLE = True
    except ImportError:
        MSS_AVAILABLE = False
        logger.warning("mss library not available, falling back to ImageGrab for Linux")
else:
    MSS_AVAILABLE = False

class CaptureMode(Enum):
    SINGLE = "single"
    MULTIPLE = "multiple"
    SCREENS = "screens"

class Capture:
    def __call__(self, mode: CaptureMode, area: Optional[tuple] = None):
        if sys.platform == "darwin":
            # macOS: Use ImageGrab (works well on macOS)
            if mode == CaptureMode.SINGLE:
                return ImageGrab.grab()
            elif mode == CaptureMode.MULTIPLE:
                return ImageGrab.grab(bbox=area)
            elif mode == CaptureMode.SCREENS:
                return ImageGrab.grab(all_screens=True)
        else:
            # Linux: Use mss for better X11 connection management
            if MSS_AVAILABLE:
                with mss() as sct:
                    if mode == CaptureMode.SINGLE:
                        # Capture primary monitor
                        monitor = sct.monitors[1]  # Skip index 0 (all monitors combined)
                        shot = sct.grab(monitor)
                        return Image.frombytes("RGB", shot.size, shot.rgb)
                    elif mode == CaptureMode.MULTIPLE:
                        # Capture specific area
                        bbox = {
                            "left": area[0],
                            "top": area[1],
                            "width": area[2] - area[0],
                            "height": area[3] - area[1]
                        }
                        shot = sct.grab(bbox)
                        return Image.frombytes("RGB", shot.size, shot.rgb)
                    elif mode == CaptureMode.SCREENS:
                        # For SCREENS mode, use MULTIPLE per monitor (handled in screenshot method)
                        raise NotImplementedError("SCREENS mode should use MULTIPLE per monitor")
            else:
                # Fallback to ImageGrab if mss not available
                if mode == CaptureMode.SINGLE:
                    return ImageGrab.grab()
                elif mode == CaptureMode.MULTIPLE:
                    return ImageGrab.grab(bbox=area)
                elif mode == CaptureMode.SCREENS:
                    return ImageGrab.grab(all_screens=True)

def apply_filter(screenshot, blur_intensity):
    return screenshot.filter(ImageFilter.GaussianBlur(blur_intensity))

def convert_image(screenshot, format):
    if screenshot.mode == "RGBA" and format == "JPEG":
        return screenshot.convert("RGB")
    return screenshot

class ScreenShot:
    def __init__(self, dir: str, format: str = "JPEG", prefix: str = "_", quality: int = 100, blur_intensity: int = 1, ssfrequencystatus: int = 1, stop_event=None):
        self.results = []  # Instance-specific state to prevent shared mutable state
        self.format = format
        if sys.platform == "darwin":
            screenshot_dir = Path.home() / "Library/Application Support/Workstatus/screenshots"
            try:
                screenshot_dir.mkdir(parents=True, exist_ok=True)
            except Exception as e:
                logger.error("Failed to create screenshot directory: %s", e)
                raise
            self.dir = str(screenshot_dir)
        else:
            self.dir = dir
        self.prefix = prefix
        self.quality = quality
        self.blur_intensity = blur_intensity
        self.ssfrequencystatus = ssfrequencystatus
        self.stop_event = stop_event

    def screenshot(self):
        screenshots = []
        monitors = get_monitors()

        if sys.platform == "darwin":
            # macOS: Use ImageGrab via Capture class
            capture = Capture()
            for monitor in monitors:
                bbox = (monitor.x, monitor.y, monitor.x + monitor.width, monitor.y + monitor.height)
                try:
                    shot = capture(mode=CaptureMode.MULTIPLE, area=bbox)
                    screenshots.append((shot, monitor))
                except OSError as e:
                    logger.error("Error capturing monitor %s: %s", monitor, e)
        else:
            # Linux: Use mss for better X11 connection management (reuses single connection)
            if MSS_AVAILABLE:
                with mss() as sct:
                    for i, monitor in enumerate(monitors):
                        logger.debug("Linux screen %s: x=%s, y=%s, width=%s, height=%s, name=%s", i + 1, monitor.x, monitor.y, monitor.width, monitor.height, getattr(monitor, 'name', 'Unknown'))
                        bbox = {
                            "left": monitor.x,
                            "top": monitor.y,
                            "width": monitor.width,
                            "height": monitor.height
                        }
                        try:
                            shot = sct.grab(bbox)
                            img = Image.frombytes("RGB", shot.size, shot.rgb)
                            screenshots.append((img, monitor))
                        except Exception as e:
                            logger.error("Error capturing monitor %s: %s", getattr(monitor, 'name', f'Monitor {i+1}'), e)
                            # FIXED: Cleanup image if it was created but append failed
                            if 'img' in locals():
                                try:
                                    img.close()
                                except Exception:
                                    pass
            else:
                # Fallback to ImageGrab if mss not available
                capture = Capture()
                for monitor in monitors:
                    bbox = (monitor.x, monitor.y, monitor.x + monitor.width, monitor.y + monitor.height)
                    try:
                        shot = capture(mode=CaptureMode.MULTIPLE, area=bbox)
                        screenshots.append((shot, monitor))
                    except OSError as e:
                        logger.error("Error capturing monitor %s: %s", monitor, e)

        return screenshots  # [(Image, monitor), ...]

    def save(self):
        results = []
        
        # First, get the number of monitors to calculate total target
        initial_screenshots = self.screenshot()
        num_monitors = len(initial_screenshots)
        # Close the initial screenshots as we just needed the count
        for screenshot, _ in initial_screenshots:
            try:
                screenshot.close()
            except Exception:
                pass
        
        # Calculate total target: ssfrequencystatus screenshots per monitor
        target_count = self.ssfrequencystatus * num_monitors
        screenshot_count = 0
        # Track screenshots per monitor to ensure each gets exactly ssfrequencystatus
        monitor_counts = {i: 0 for i in range(num_monitors)}
        
        logger.info("[SCREENSHOT] Capture plan: %s monitor(s), %s screenshot(s) per monitor = %s total screenshot(s)", num_monitors, self.ssfrequencystatus, target_count)

        # Capture screenshots until each monitor has ssfrequencystatus screenshots
        while screenshot_count < target_count:
            if self.stop_event and self.stop_event.is_set():
                break
            screenshots = self.screenshot()
            last_processed_idx = -1

            for i, (screenshot, monitor) in enumerate(screenshots):
                if self.stop_event and self.stop_event.is_set():
                    # Close remaining screenshots before breaking
                    for remaining_idx in range(i, len(screenshots)):
                        try:
                            screenshots[remaining_idx][0].close()
                        except Exception:
                            pass
                    break
                
                # Check if this monitor has already reached its target
                if monitor_counts[i] >= self.ssfrequencystatus:
                    # This monitor already has enough screenshots, close and skip
                    logger.debug("[SCREENSHOT] Monitor %s already has %s/%s screenshots, skipping", i+1, monitor_counts[i], self.ssfrequencystatus)
                    screenshot.close()
                    continue
                
                if screenshot_count >= target_count:
                    # Close unused screenshots if we've reached total target
                    screenshot.close()
                    last_processed_idx = i
                    break
                
                original_screenshot = screenshot  # Keep reference for cleanup
                processed_screenshot = None
                
                try:
                    if self.blur_intensity == 1:
                        blurred = apply_filter(screenshot, 5)
                        # FIXED: Close original if new image created
                        if blurred is not screenshot:
                            screenshot.close()
                        screenshot = blurred
                        processed_screenshot = screenshot

                    converted = convert_image(screenshot, self.format)
                    # FIXED: Close previous if new image created
                    if converted is not screenshot:
                        screenshot.close()
                    screenshot = converted
                    processed_screenshot = screenshot

                    filename = make_screenshot_filename(
                        self.format, offset_seconds=screenshot_count
                    )
                    path = os.path.join(self.dir, f"{self.prefix}{filename}")

                    screenshot.save(path, format=self.format, quality=self.quality)
                    results.append((screenshot, path, filename))
                    screenshot_count += 1
                    monitor_counts[i] += 1
                    last_processed_idx = i
                    logger.debug("[SCREENSHOT] Saved screenshot %s/%s for monitor %s (monitor %s has %s/%s)", screenshot_count, target_count, i+1, i+1, monitor_counts[i], self.ssfrequencystatus)
                except Exception as e:
                    # FIXED: Ensure cleanup on error
                    logger.error("Error processing screenshot: %s", e)
                    if processed_screenshot is not None and processed_screenshot is not original_screenshot:
                        try:
                            processed_screenshot.close()
                        except Exception:
                            pass
                    if original_screenshot is not None and original_screenshot is not processed_screenshot:
                        try:
                            original_screenshot.close()
                        except Exception:
                            pass
                    raise
            
            # Close any remaining unused screenshots from this capture
            if screenshot_count >= target_count and last_processed_idx >= 0:
                for remaining_idx in range(last_processed_idx + 1, len(screenshots)):
                    if remaining_idx < len(screenshots):
                        try:
                            screenshots[remaining_idx][0].close()
                        except Exception:
                            pass
        
        per_monitor = ", ".join(f"Monitor {i+1}: {monitor_counts.get(i, 0)}" for i in range(num_monitors))
        logger.info("[SCREENSHOT] Capture complete: %s screenshot(s). Per monitor: %s", screenshot_count, per_monitor)
        return results

    def __enter__(self):
        if not os.path.exists(self.dir):
            os.makedirs(self.dir)
        self.results = self.save()
        return self.results

    def __exit__(self, type, value, traceback):
        for result in self.results:
            if result[0] is not None:
                result[0].close()
