"""
Screen capture functionality for FMUS-VID.

This module provides tools for capturing screenshots of the screen or specific regions.
"""

import os
import platform
from pathlib import Path
from typing import Optional, Tuple, Union, List
import tempfile
import time

import numpy as np
from PIL import Image, ImageGrab
import cv2  # Added for image processing

# Define a simple Frame class for the capture module instead of importing from core
class Frame:
    """Simplified Frame class for captured screenshots."""

    def __init__(self, frame_data):
        """Initialize a Frame with numpy array data."""
        self.data = frame_data

    def save(self, path):
        """Save the frame to a file."""
        path = Path(path)
        cv2.imwrite(str(path), cv2.cvtColor(self.data, cv2.COLOR_RGB2BGR))
        return path

    def to_numpy(self):
        """Return the frame as a numpy array."""
        return self.data

    def to_pil(self):
        """Convert to PIL Image."""
        return Image.fromarray(self.data)


def capture_screen(output_path: Optional[Union[str, Path]] = None) -> Union[Frame, Path]:
    """
    Capture a screenshot of the entire screen.

    Args:
        output_path: Optional path to save the screenshot. If None, returns a Frame object.

    Returns:
        Frame object if output_path is None, otherwise the Path to the saved image.

    Example:
        >>> # Capture and get Frame object
        >>> frame = fmusvid.capture.capture_screen()
        >>> # Capture and save to file
        >>> path = fmusvid.capture.capture_screen("screenshot.png")
    """
    # Capture the screen
    screenshot = ImageGrab.grab()

    # Convert to numpy array for processing
    frame_array = np.array(screenshot)

    # Create a Frame object
    frame = Frame(frame_array)

    # If output path is provided, save the image
    if output_path:
        output_path = Path(output_path)
        return frame.save(output_path)

    return frame


def capture_region(region: Tuple[int, int, int, int],
                 output_path: Optional[Union[str, Path]] = None) -> Union[Frame, Path]:
    """
    Capture a screenshot of a specific region of the screen.

    Args:
        region: Tuple of (left, top, right, bottom) coordinates
        output_path: Optional path to save the screenshot. If None, returns a Frame object.

    Returns:
        Frame object if output_path is None, otherwise the Path to the saved image.

    Example:
        >>> # Capture region (100, 100, 500, 400) and get Frame object
        >>> frame = fmusvid.capture.capture_region((100, 100, 500, 400))
        >>> # Capture region and save to file
        >>> path = fmusvid.capture.capture_region((100, 100, 500, 400), "region.png")
    """
    # Validate region coordinates
    left, top, right, bottom = region
    if left >= right or top >= bottom:
        raise ValueError("Invalid region coordinates. Ensure left < right and top < bottom.")

    # Capture the region
    screenshot = ImageGrab.grab(bbox=region)

    # Convert to numpy array for processing
    frame_array = np.array(screenshot)

    # Create a Frame object
    frame = Frame(frame_array)

    # If output path is provided, save the image
    if output_path:
        output_path = Path(output_path)
        return frame.save(output_path)

    return frame


def select_region() -> Tuple[int, int, int, int]:
    """
    Interactive tool to select a region of the screen.

    Returns:
        Tuple of (left, top, right, bottom) coordinates

    Example:
        >>> # Select a region interactively
        >>> region = fmusvid.capture.select_region()
        >>> # Capture the selected region
        >>> frame = fmusvid.capture.capture_region(region)
    """
    try:
        import tkinter as tk
        from tkinter import simpledialog
    except ImportError:
        raise ImportError("Tkinter is required for interactive region selection. "
                         "Install it or specify region coordinates manually.")

    # Create a full screen transparent window
    root = tk.Tk()
    root.attributes('-alpha', 0.3)
    root.attributes('-fullscreen', True)

    # Variables to store the selection coordinates
    start_x = start_y = end_x = end_y = 0
    selection_made = False

    # Create a canvas
    canvas = tk.Canvas(root, cursor="cross")
    canvas.pack(fill=tk.BOTH, expand=True)

    # Rectangle ID for the selection
    rect_id = None

    def on_mouse_down(event):
        nonlocal start_x, start_y, rect_id
        start_x, start_y = event.x, event.y
        # Create a rectangle if it doesn't exist
        if rect_id:
            canvas.delete(rect_id)
        rect_id = canvas.create_rectangle(start_x, start_y, start_x, start_y,
                                         outline='red', width=2)

    def on_mouse_move(event):
        nonlocal rect_id
        if rect_id:
            canvas.coords(rect_id, start_x, start_y, event.x, event.y)

    def on_mouse_up(event):
        nonlocal end_x, end_y, selection_made
        end_x, end_y = event.x, event.y
        selection_made = True
        root.quit()

    # Bind mouse events
    canvas.bind("<ButtonPress-1>", on_mouse_down)
    canvas.bind("<B1-Motion>", on_mouse_move)
    canvas.bind("<ButtonRelease-1>", on_mouse_up)

    # Run the main loop
    root.mainloop()

    # Clean up
    root.destroy()

    # Sort coordinates so left < right and top < bottom
    left = min(start_x, end_x)
    top = min(start_y, end_y)
    right = max(start_x, end_x)
    bottom = max(start_y, end_y)

    return (left, top, right, bottom)
