
import cv2
import numpy as np
from PIL import Image, ImageGrab
from pynput.mouse import Button, Controller
import platform
import ctypes
from . import dpi_manager
from screeninfo import get_monitors


# Initialize the controller once to save performance
_mouse_controller = Controller()
dpi_manager.enable_dpi_awareness()
# --- Screen Routing & Configuration ---
# Maps Logical Screen Index (Script) -> Physical Screen Index (Hardware)
_SCREEN_ROUTER = {}


def route_screen(logical_screen, physical_screen):
    """
    Redirects searches intended for 'logical_screen' to 'physical_screen'.
    Useful when a script written for Screen 1 needs to run on Screen 0.
    Example: route_screen(source=1, target=0)
    """
    _SCREEN_ROUTER[logical_screen] = physical_screen


def _resolve_screen(screen_idx):
    """Resolves logical screen index to physical index."""
    return _SCREEN_ROUTER.get(screen_idx, screen_idx)


def get_monitors_safe():
    """
    Returns a list of bounding boxes (x, y, w, h) for all connected monitors.
    THIS IS THE SINGLE SOURCE OF TRUTH for both the GUI and the Runtime.
    """
    monitors = []

    # Priority 1: Use screeninfo (Recommended)
    if get_monitors:
        try:
            # We trust the order returned by screeninfo implicitly.
            si_monitors = get_monitors()
            for m in si_monitors:
                monitors.append((m.x, m.y, m.width, m.height))
            return monitors
        except Exception as e:
            print(f"screeninfo error: {e}")

    # Priority 2: Windows Fallback (ctypes)
    if not monitors and platform.system() == "Windows":
        try:
            user32 = ctypes.windll.user32

            class RECT(ctypes.Structure):
                _fields_ = [("left", ctypes.c_long), ("top", ctypes.c_long),
                            ("right", ctypes.c_long), ("bottom", ctypes.c_long)]

            MONITORENUMPROC = ctypes.WINFUNCTYPE(ctypes.c_bool, ctypes.c_void_p, ctypes.c_void_p, ctypes.POINTER(RECT),
                                                 ctypes.c_double)

            def _monitor_enum_proc(hMonitor, hdcMonitor, lprcMonitor, dwData):
                r = lprcMonitor.contents
                monitors.append((r.left, r.top, r.right - r.left, r.bottom - r.top))
                return True

            user32.EnumDisplayMonitors(None, None, MONITORENUMPROC(_monitor_enum_proc), 0)
        except Exception:
            pass

    # Priority 3: Last Resort (Virtual Desktop)
    if not monitors:
        try:
            img = ImageGrab.grab()
            monitors.append((0, 0, img.width, img.height))
        except Exception:
            monitors.append((0, 0, 1920, 1080))

    return monitors


def _resize_template(needle_pil, scale_factor):
    """Resizes the needle image by the scale factor using Lanczos resampling."""
    if scale_factor == 1.0:
        return needle_pil

    w, h = needle_pil.size
    new_w = int(max(1, w * scale_factor))
    new_h = int(max(1, h * scale_factor))
    return needle_pil.resize((new_w, new_h), Image.Resampling.LANCZOS)


def _load_image(img):
    """Helper to load image from path or PIL Image."""
    if isinstance(img, str):
        return Image.open(img)
    return img


def _non_max_suppression(boxes, overlap_thresh):
    """
    Standard Non-Maximum Suppression (NMS) to remove overlapping bounding boxes.
    """
    if len(boxes) == 0:
        return []

    boxes = np.array(boxes, dtype=np.float32)

    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 0] + boxes[:, 2]
    y2 = boxes[:, 1] + boxes[:, 3]

    area = (x2 - x1 + 1) * (y2 - y1 + 1)
    idxs = np.argsort(y2)

    pick = []

    while len(idxs) > 0:
        last = len(idxs) - 1
        i = idxs[last]
        pick.append(i)

        xx1 = np.maximum(x1[i], x1[idxs[:last]])
        yy1 = np.maximum(y1[i], y1[idxs[:last]])
        xx2 = np.minimum(x2[i], x2[idxs[:last]])
        yy2 = np.minimum(y2[i], y2[idxs[:last]])

        w = np.maximum(0, xx2 - xx1 + 1)
        h = np.maximum(0, yy2 - yy1 + 1)

        overlap = (w * h) / area[idxs[:last]]

        idxs = np.delete(idxs, np.concatenate(([last],
                                               np.where(overlap > overlap_thresh)[0])))

    return boxes[pick].astype("int").tolist()


def _run_template_match(needleImage, haystackImage, grayscale=False):
    """
    Shared logic for preparing images and running cv2.matchTemplate.
    """
    haystack_pil = _load_image(haystackImage)
    haystack_np = np.array(haystack_pil)

    if haystack_pil.mode == 'RGB':
        haystack = cv2.cvtColor(haystack_np, cv2.COLOR_RGB2BGR)
    elif haystack_pil.mode == 'RGBA':
        haystack = cv2.cvtColor(haystack_np, cv2.COLOR_RGBA2BGR)
    else:
        haystack = haystack_np
        if len(haystack.shape) == 2:
            haystack = cv2.cvtColor(haystack, cv2.COLOR_GRAY2BGR)

    needle_pil = _load_image(needleImage)
    needle_np = np.array(needle_pil)

    needle = None
    mask = None

    if needle_pil.mode == 'RGBA':
        needle_bgra = cv2.cvtColor(needle_np, cv2.COLOR_RGBA2BGRA)
        needle = needle_bgra[:, :, :3]
        mask = needle_bgra[:, :, 3]
    else:
        if needle_pil.mode == 'RGB':
            needle = cv2.cvtColor(needle_np, cv2.COLOR_RGB2BGR)
        else:
            needle = needle_np
            if len(needle.shape) == 2:
                needle = cv2.cvtColor(needle, cv2.COLOR_GRAY2BGR)

    if grayscale:
        haystack = cv2.cvtColor(haystack, cv2.COLOR_BGR2GRAY)
        needle = cv2.cvtColor(needle, cv2.COLOR_BGR2GRAY)

    if mask is not None and not grayscale:
        res = cv2.matchTemplate(haystack, needle, cv2.TM_CCOEFF_NORMED, mask=mask)
    else:
        res = cv2.matchTemplate(haystack, needle, cv2.TM_CCOEFF_NORMED)

    h, w = needle.shape[:2]
    return res, w, h


def locate(needleImage, haystackImage, grayscale=False, confidence=0.9):
    """
    Locate the BEST instance of 'needleImage' inside 'haystackImage'.
    """
    res, w, h = _run_template_match(needleImage, haystackImage, grayscale)
    minVal, maxVal, minLoc, maxLoc = cv2.minMaxLoc(res)

    if maxVal >= confidence:
        return (maxLoc[0], maxLoc[1], w, h)

    return None


def locateAll(needleImage, haystackImage, grayscale=False, confidence=0.9, overlap_threshold=0.5):
    """
    Locate all instances of 'needleImage' inside 'haystackImage'.
    """
    res, w, h = _run_template_match(needleImage, haystackImage, grayscale)

    loc = np.where(res >= confidence)
    rects = []
    for pt in zip(*loc[::-1]):
        rects.append([int(pt[0]), int(pt[1]), int(w), int(h)])

    if overlap_threshold < 1.0 and len(rects) > 1:
        rects = _non_max_suppression(rects, overlap_threshold)

    rects.sort(key=lambda r: (r[1], r[0]))

    return [tuple(r) for r in rects]


def locateAllOnScreen(image, region=None, screen=0, grayscale=False, confidence=0.9, overlap_threshold=0.5,
                      original_resolution=None):
    """
    Locate all instances of 'image' on the screen.
    """
    haystack_pil, offset_x, offset_y, scale_factor = _prepare_screen_capture(region, screen, original_resolution)
    needle_pil = _load_image(image)
    if scale_factor != 1.0:
        needle_pil = _resize_template(needle_pil, scale_factor)

    rects = locateAll(needle_pil, haystack_pil, grayscale, confidence, overlap_threshold)

    if offset_x or offset_y:
        final_rects = []
        for (x, y, w, h) in rects:
            final_rects.append((x + offset_x, y + offset_y, w, h))
        return final_rects

    return rects


def locateOnScreen(image, region=None, screen=0, grayscale=False, confidence=0.9,
                   original_resolution=None):
    """
    Locate the first instance of 'image' on the screen.
    """
    # 1. Prepare screen and image (Same as your existing functions)
    haystack_pil, offset_x, offset_y, scale_factor = _prepare_screen_capture(region, screen, original_resolution)
    needle_pil = _load_image(image)
    if scale_factor != 1.0:
        needle_pil = _resize_template(needle_pil, scale_factor)

    # 2. Get raw matches (Generator or List)
    # We use locateAll because it returns results sorted by position (top-left to bottom-right)
    # whereas locate() returns results sorted by confidence.
    rects = locateAll(needle_pil, haystack_pil, grayscale, confidence)

    # 3. OPTIMIZATION: Extract only the first item immediately
    # This works whether 'rects' is a list or a generator
    try:
        first_match = next(iter(rects))
    except StopIteration:
        return None  # No matches found

    # 4. Apply offset ONLY to the first match
    # (Your original function applied this math to every single match found)
    x, y, w, h = first_match
    return (x + offset_x, y + offset_y, w, h)


def _prepare_screen_capture(region, screen_idx, original_resolution):
    """
    Internal helper that calls get_monitors_safe() to ensure alignment.
    """
    # 1. Routing
    physical_screen = _resolve_screen(screen_idx)

    # CALL THE PUBLIC SAFE FUNCTION
    monitors = get_monitors_safe()

    # Safety Fallback
    if physical_screen >= len(monitors):
        print(f"Warning: Screen {physical_screen} not found. Falling back to Primary (0).")
        physical_screen = 0

    target_monitor_rect = monitors[physical_screen]

    # 2. Determine Capture Area
    offset_x, offset_y = 0, 0

    if region:
        x, y, w, h = region
        haystack_pil = ImageGrab.grab(bbox=(x, y, x + w, y + h))
        offset_x, offset_y = x, y
    else:
        x, y, w, h = target_monitor_rect
        haystack_pil = ImageGrab.grab(bbox=(x, y, x + w, y + h))
        offset_x, offset_y = x, y

    # 3. Determine Scale Factor
    scale_factor = 1.0

    if original_resolution:
        orig_w, orig_h = original_resolution
        target_monitor_h = target_monitor_rect[3]
        scale_factor = target_monitor_h / float(orig_h)
        if abs(scale_factor - 1.0) < 0.02:
            scale_factor = 1.0

    return haystack_pil, offset_x, offset_y, scale_factor


def clickimage(match, offset=(0, 0), button='left', clicks=1):
    """
    Clicks a location with an optional offset using pynput.
    """
    if not match:
        print("Debug: No match found, skipping click.")
        return

    x, y, w, h = match
    center_x = x + (w / 2)
    center_y = y + (h / 2)
    target_x = center_x + offset[0]
    target_y = center_y + offset[1]

    _mouse_controller.position = (target_x, target_y)

    pynput_button = Button.left
    if button == 'right':
        pynput_button = Button.right
    elif button == 'middle':
        pynput_button = Button.middle

    _mouse_controller.click(pynput_button, clicks)