import subprocess
import base64
from pathlib import Path
from PIL import ImageGrab
from uuid import uuid4
from screeninfo import get_monitors
import platform
if platform.system() == "Darwin":
    import Quartz  # uncomment this line if you are on macOS
    
from PIL import ImageGrab
from functools import partial
from .base import BaseAnthropicTool, ToolError, ToolResult


OUTPUT_DIR = "./tmp/outputs"

def get_screenshot(selected_screen: int = 0, resize: bool = True, target_width: int = 1920, target_height: int = 1080):
    # print(f"get_screenshot selected_screen: {selected_screen}")
    
    # Get screen width and height using Windows command
    display_num = None
    offset_x = 0
    offset_y = 0
    selected_screen = selected_screen   
    width, height = _get_screen_size()    

    """Take a screenshot of the current screen and return a ToolResult with the base64 encoded image."""
    output_dir = Path(OUTPUT_DIR)
    output_dir.mkdir(parents=True, exist_ok=True)
    path = output_dir / f"screenshot_{uuid4().hex}.png"

    ImageGrab.grab = partial(ImageGrab.grab, all_screens=True)

    # Detect platform
    system = platform.system()

    if system == "Windows":
        # Windows: Use screeninfo to get monitor details
        screens = get_monitors()

        # Sort screens by x position to arrange from left to right
        sorted_screens = sorted(screens, key=lambda s: s.x)

        if selected_screen < 0 or selected_screen >= len(screens):
            raise IndexError("Invalid screen index.")

        screen = sorted_screens[selected_screen]
        bbox = (screen.x, screen.y, screen.x + screen.width, screen.y + screen.height)

    elif system == "Darwin":  # macOS
        # macOS: Use Quartz to get monitor details
        max_displays = 32  # Maximum number of displays to handle
        active_displays = Quartz.CGGetActiveDisplayList(max_displays, None, None)[1]

        # Get the display bounds (resolution) for each active display
        screens = []
        for display_id in active_displays:
            bounds = Quartz.CGDisplayBounds(display_id)
            screens.append({
                'id': display_id,
                'x': int(bounds.origin.x),
                'y': int(bounds.origin.y),
                'width': int(bounds.size.width),
                'height': int(bounds.size.height),
                'is_primary': Quartz.CGDisplayIsMain(display_id)  # Check if this is the primary display
            })

        # Sort screens by x position to arrange from left to right
        sorted_screens = sorted(screens, key=lambda s: s['x'])
        # print(f"Darwin sorted_screens: {sorted_screens}")

        if selected_screen < 0 or selected_screen >= len(screens):
            raise IndexError("Invalid screen index.")

        screen = sorted_screens[selected_screen]
        
        bbox = (screen['x'], screen['y'], screen['x'] + screen['width'], screen['y'] + screen['height'])

    else:  # Linux or other OS
        cmd = "xrandr | grep ' primary' | awk '{print $4}'"
        try:
            output = subprocess.check_output(cmd, shell=True).decode()
            resolution = output.strip().split()[0]
            width, height = map(int, resolution.split('x'))
            bbox = (0, 0, width, height)  # Assuming single primary screen for simplicity
        except subprocess.CalledProcessError:
            raise RuntimeError("Failed to get screen resolution on Linux.")

    # Take screenshot using the bounding box
    screenshot = ImageGrab.grab(bbox=bbox)

    # Set offsets (for potential future use)
    offset_x = screen['x'] if system == "Darwin" else screen.x
    offset_y = screen['y'] if system == "Darwin" else screen.y

    # # Resize if 
    if resize:
        screenshot = screenshot.resize((target_width, target_height))

    # Save the screenshot
    screenshot.save(str(path))

    if path.exists():
        # Return a ToolResult instance instead of a dictionary
        return screenshot, str(path)
    
    raise ToolError(f"Failed to take screenshot: {path} does not exist.")
    
    


def _get_screen_size(selected_screen: int = 0):
    if platform.system() == "Windows":
        # Use screeninfo to get primary monitor on Windows
        screens = get_monitors()

        # Sort screens by x position to arrange from left to right
        sorted_screens = sorted(screens, key=lambda s: s.x)
        if selected_screen is None:
            primary_monitor = next((m for m in get_monitors() if m.is_primary), None)
            return primary_monitor.width, primary_monitor.height
        elif selected_screen < 0 or selected_screen >= len(screens):
            raise IndexError("Invalid screen index.")
        else:
            screen = sorted_screens[selected_screen]
            return screen.width, screen.height
    elif platform.system() == "Darwin":
        # macOS part using Quartz to get screen information
        max_displays = 32  # Maximum number of displays to handle
        active_displays = Quartz.CGGetActiveDisplayList(max_displays, None, None)[1]

        # Get the display bounds (resolution) for each active display
        screens = []
        for display_id in active_displays:
            bounds = Quartz.CGDisplayBounds(display_id)
            screens.append({
                'id': display_id,
                'x': int(bounds.origin.x),
                'y': int(bounds.origin.y),
                'width': int(bounds.size.width),
                'height': int(bounds.size.height),
                'is_primary': Quartz.CGDisplayIsMain(display_id)  # Check if this is the primary display
            })

        # Sort screens by x position to arrange from left to right
        sorted_screens = sorted(screens, key=lambda s: s['x'])

        if selected_screen is None:
            # Find the primary monitor
            primary_monitor = next((screen for screen in screens if screen['is_primary']), None)
            if primary_monitor:
                return primary_monitor['width'], primary_monitor['height']
            else:
                raise RuntimeError("No primary monitor found.")
        elif selected_screen < 0 or selected_screen >= len(screens):
            raise IndexError("Invalid screen index.")
        else:
            # Return the resolution of the selected screen
            screen = sorted_screens[selected_screen]
            return screen['width'], screen['height']

    else:  # Linux or other OS
        cmd = "xrandr | grep ' primary' | awk '{print $4}'"
        try:
            output = subprocess.check_output(cmd, shell=True).decode()
            resolution = output.strip().split()[0]
            width, height = map(int, resolution.split('x'))
            return width, height
        except subprocess.CalledProcessError:
            raise RuntimeError("Failed to get screen resolution on Linux.")
