"""
Capture a screenshot of the full screen or a specific window by title.

Usage:
    python capture.py --output screenshot.png
    python capture.py --window "Chrome" --output chrome.png
    python capture.py --window "Code" --activate --delay 1 --output vscode.png
"""
import argparse
import json
import sys
import time

try:
    import mss
    import mss.tools
except ImportError:
    print("Error: mss is not installed. Run the setup script first.", file=sys.stderr)
    sys.exit(1)

try:
    from PIL import Image
except ImportError:
    print("Error: Pillow is not installed. Run the setup script first.", file=sys.stderr)
    sys.exit(1)


def find_window(title_substring):
    """Find a window whose title contains the given substring (case-insensitive)."""
    try:
        import pygetwindow as gw
    except ImportError:
        print("Error: pygetwindow is not installed. Run the setup script first.", file=sys.stderr)
        sys.exit(1)

    title_lower = title_substring.lower()
    matches = []
    all_windows = []

    for w in gw.getAllWindows():
        if not w.title.strip() or w.width <= 0 or w.height <= 0:
            continue
        all_windows.append(w.title)
        if title_lower in w.title.lower():
            matches.append(w)

    if not matches:
        print(f"Error: No window found matching '{title_substring}'.", file=sys.stderr)
        print(f"Available windows:", file=sys.stderr)
        for t in all_windows:
            print(f"  - {t}", file=sys.stderr)
        sys.exit(1)

    if len(matches) > 1:
        print(f"Warning: Multiple windows match '{title_substring}'. Using the first:", file=sys.stderr)
        for m in matches:
            print(f"  - {m.title}", file=sys.stderr)

    return matches[0]


def activate_window(window):
    """Bring a window to the foreground."""
    try:
        import win32gui
        import win32con
        hwnd = window._hWnd
        # Restore if minimized
        if win32gui.IsIconic(hwnd):
            win32gui.ShowWindow(hwnd, win32con.SW_RESTORE)
        win32gui.SetForegroundWindow(hwnd)
    except Exception as e:
        print(f"Warning: Could not activate window: {e}", file=sys.stderr)


def capture_full_screen(output_path):
    """Capture the full primary monitor."""
    with mss.mss() as sct:
        monitor = sct.monitors[1]  # Primary monitor
        screenshot = sct.grab(monitor)
        img = Image.frombytes("RGB", screenshot.size, screenshot.bgra, "raw", "BGRX")
        img.save(output_path, "PNG")
    print(json.dumps({"status": "ok", "output": output_path, "width": img.width, "height": img.height}))


def capture_window(window, output_path):
    """Capture a specific window region."""
    region = {
        "left": max(0, window.left),
        "top": max(0, window.top),
        "width": window.width,
        "height": window.height,
    }

    with mss.mss() as sct:
        screenshot = sct.grab(region)
        img = Image.frombytes("RGB", screenshot.size, screenshot.bgra, "raw", "BGRX")
        img.save(output_path, "PNG")
    print(json.dumps({"status": "ok", "output": output_path, "width": img.width, "height": img.height, "window": window.title}))


def main():
    parser = argparse.ArgumentParser(description="Capture a screenshot.")
    parser.add_argument("--output", required=True, help="Output PNG file path")
    parser.add_argument("--window", default=None, help="Window title substring to match")
    parser.add_argument("--delay", type=float, default=0, help="Seconds to wait before capture")
    parser.add_argument("--activate", action="store_true", help="Bring window to foreground first")
    args = parser.parse_args()

    if args.window:
        window = find_window(args.window)
        if args.activate:
            activate_window(window)
        if args.delay > 0:
            time.sleep(args.delay)
        capture_window(window, args.output)
    else:
        if args.delay > 0:
            time.sleep(args.delay)
        capture_full_screen(args.output)


if __name__ == "__main__":
    main()