#!/usr/bin/env python3
"""
Full-page screenshot (v2): Scroll-and-stitch the CURRENT window without WebDriver.

This module exposes one public function:
  - save_fullpage_screenshot_v2(...)

And a minimal CLI:
  python fullpage_screenshot_v2.py OUT.png [--window-title-contains "Chrome"] [--activate] [--maximize]
  python fullpage_screenshot_v2.py "xero_fullpage.png" --use-pagedown --window-title-contains "Chrome" --activate --maximize

Dependencies: pyautogui, pygetwindow, Pillow (installed in requirements.txt)
"""

from __future__ import annotations

from typing import Optional

from PIL import Image
import os
import subprocess
import sys


class _SimpleWindow:
    """Minimal window wrapper for Linux/X11 when pygetwindow is unavailable.

    Provides left/top/width/height attributes and no-op activate/maximize methods.
    When possible, activate/maximize will use xdotool/wmctrl on the stored window id.
    """

    def __init__(self, left: int, top: int, width: int, height: int, wid: str | None = None):
        self.left = int(left)
        self.top = int(top)
        self.width = int(width)
        self.height = int(height)
        self._wid = wid

    def activate(self) -> None:
        if not self._wid:
            return
        try:
            subprocess.run(["xdotool", "windowactivate", "--sync", str(self._wid)], check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        except Exception:
            pass

    def maximize(self) -> None:
        if not self._wid:
            return
        try:
            # Best-effort maximize via wmctrl
            subprocess.run(["wmctrl", "-ir", str(self._wid), "-b", "add,maximized_vert,maximized_horz"], check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        except Exception:
            pass


def _locate_window_linux_x11(window_title_contains: str | None):
    """Locate the active window geometry on Linux/X11 via xdotool/wmctrl.

    If window_title_contains is provided, best-effort attempt to focus a matching
    Chromium window before reading geometry. Requires a running X session with DISPLAY.
    """
    # Ensure we have access to an X11 session
    disp = os.environ.get("DISPLAY", "")
    if not disp:
        raise RuntimeError("DISPLAY is not set; cannot query active window under X11")

    # Optionally try to activate a Chromium window if requested
    if window_title_contains:
        try:
            # Try wmctrl by class first (covers Chrome/Chromium/Edge)
            for cls in ("google-chrome.Google-chrome", "chromium.Chromium", "microsoft-edge.Microsoft-edge"):
                subprocess.run(["wmctrl", "-x", "-a", cls], check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            # Fallback: generic name match
            subprocess.run(["wmctrl", "-a", str(window_title_contains)], check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        except Exception:
            pass

    # Get active window id
    try:
        wid = subprocess.check_output(["xdotool", "getactivewindow"], text=True).strip()
    except Exception as exc:
        raise RuntimeError("xdotool not available or cannot access X session") from exc

    # Read geometry in shell format and parse
    try:
        geo_out = subprocess.check_output(["xdotool", "getwindowgeometry", "--shell", wid], text=True)
        # Example lines: X=123\nY=45\nWIDTH=800\nHEIGHT=600
        kv = {}
        for line in geo_out.splitlines():
            if "=" in line:
                k, v = line.split("=", 1)
                kv[k.strip()] = v.strip()
        left = int(kv.get("X", "0"))
        top = int(kv.get("Y", "0"))
        width = int(kv.get("WIDTH", "0"))
        height = int(kv.get("HEIGHT", "0"))
        if width <= 0 or height <= 0:
            raise RuntimeError("Active window has invalid geometry via xdotool")
        return _SimpleWindow(left, top, width, height, wid=wid)
    except Exception as exc:
        raise RuntimeError("Failed to read active window geometry via xdotool") from exc


def _locate_window(window_title_contains: str | None):
    try:
        import pygetwindow as gw
    except Exception as exc:
        # On Linux/X11, fall back to xdotool/wmctrl-based geometry
        if sys.platform != "win32":
            try:
                return _locate_window_linux_x11(window_title_contains)
            except Exception as exc2:
                raise RuntimeError("pygetwindow unavailable and Linux/X11 fallback failed") from exc2
        raise RuntimeError("pygetwindow is required for v2 capture") from exc

    target = None
    if window_title_contains:
        needle = window_title_contains.lower()
        for w in gw.getAllWindows():
            title = (getattr(w, 'title', '') or '').lower()
            if needle in title:
                target = w
                break
    if not target:
        target = getattr(gw, 'getActiveWindow', None)() if hasattr(gw, 'getActiveWindow') else None
    if not target:
        raise RuntimeError("No active window detected")
    return target


def _focus_content_region(x: int | None = None, y: int | None = None) -> None:
    """Bring focus to the page by clicking at (x,y) when provided, else screen center.

    Prefer passing window content coordinates from callers to avoid clicking other apps.
    """
    try:
        import pyautogui as _pg
        import time as _time
        if x is None or y is None:
            try:
                sw, sh = _pg.size()
            except Exception:
                sw, sh = 1920, 1080
            x = int(max(10, sw // 2))
            y = int(max(10, sh // 2))
        _pg.moveTo(int(x), int(y))
        _pg.click(int(x), int(y))
        _time.sleep(0.1)
    except Exception:
        pass

def save_fullpage_screenshot_v2(
    output_path: str,
    window_title_contains: Optional[str] = None,
    *,
    activate: bool = True,
    maximize: bool = False,
    delay_seconds: float | None = 0.5,
    content_top_offset: int = 0,
    content_bottom_offset: int = 0,
    overlap_px: int = 80,
    scroll_px: int = 700,
    scroll_pause: float = 1.0,
    max_shots: int = 25,
    use_pagedown: bool = False,
) -> None:
    """
    Capture a full-height screenshot of the CURRENT window by scrolling + stitching.

    - window_title_contains: optional title substring to target a window
    - activate/maximize: focus and maximize before capture
    - content_top_offset/content_bottom_offset: crop toolbars/taskbar from window bounds
    - overlap_px: pixel overlap between captures (for seam-free stitching)
    - scroll_px/scroll_pause/max_shots: scrolling behavior
    - use_pagedown: use PageDown instead of mouse wheel scroll
    """
    try:
        import pyautogui
        import time
    except Exception as exc:
        raise RuntimeError("pyautogui is required for v2 capture") from exc

    target = _locate_window(window_title_contains)

    # Intentionally do NOT call window activate/maximize to avoid stealing focus.
    # We rely on content-region click focusing instead.
    if delay_seconds and delay_seconds > 0:
        time.sleep(delay_seconds)

    left = int(getattr(target, 'left', 0))
    top = int(getattr(target, 'top', 0))
    width = int(getattr(target, 'width', 0))
    height = int(getattr(target, 'height', 0))
    if width <= 0 or height <= 0:
        raise RuntimeError("Active window has invalid size")

    # Define content region inside the window
    region_top = top + max(0, int(content_top_offset))
    region_height = max(1, height - int(content_top_offset) - int(content_bottom_offset))
    region = (left, region_top, width, region_height)

    # Ensure the web content has focus: click at screen center (requested behavior)
    _focus_content_region()

    # Go to top
    for _ in range(3):
        pyautogui.keyDown('ctrl')
        pyautogui.press('home')
        pyautogui.keyUp('ctrl')
        time.sleep(0.1)

    images: list[Image.Image] = []
    last_im_bytes = None
    for i in range(max_shots):
        # Try pyautogui first (works well when called directly), then mss as a fallback
        try:
            im = pyautogui.screenshot(region=region)
        except Exception:
            try:
                import mss  # type: ignore
                with mss.mss() as sct:
                    mon = {"left": int(region[0]), "top": int(region[1]), "width": int(region[2]), "height": int(region[3])}
                    sct_img = sct.grab(mon)
                    im = Image.frombytes("RGB", sct_img.size, sct_img.bgra, "raw", "BGRA")
            except Exception as _exc:
                raise RuntimeError(f"Region screenshot failed: {_exc}")
        im = im.convert('RGB')
        buf = im.tobytes()
        if last_im_bytes is not None and buf == last_im_bytes:
            break
        images.append(im)
        last_im_bytes = buf

        if i < max_shots - 1:
            if use_pagedown:
                pyautogui.press('pagedown')
            else:
                pyautogui.scroll(-int(scroll_px))
            time.sleep(float(scroll_pause))

    if not images:
        raise RuntimeError("No images captured during v2 capture")

    # Dynamic alignment with OpenCV to avoid seams when little scrolling is needed
    try:
        import numpy as np
        import cv2
        use_cv = True
    except Exception:
        use_cv = False

    # Normalize widths to the minimum width across frames to keep edges aligned
    base_w = min(img.size[0] for img in images)
    images = [img.crop((0, 0, base_w, img.size[1])) for img in images]

    # Start with the first frame
    out = Image.new('RGB', (base_w, images[0].size[1]), (255, 255, 255))
    out.paste(images[0], (0, 0))

    keep_overlap_px = 12  # safety margin to avoid cutting content at seams
    seam_bias_px = 16     # additional bias to include a bit more top from next frame
    center_strip_ratio = 0.7  # use center 70% width for alignment to avoid sticky sidebars
    for i in range(1, len(images)):
        prev = images[i - 1]
        curr = images[i]
        prev_h = prev.size[1]
        curr_w, curr_h = curr.size

        if use_cv:
            prev_gray = cv2.cvtColor(np.array(prev), cv2.COLOR_RGB2GRAY)
            curr_gray = cv2.cvtColor(np.array(curr), cv2.COLOR_RGB2GRAY)
            # Use only the center strip for matching
            cx = int(base_w * (1 - center_strip_ratio) / 2)
            cw = int(base_w * center_strip_ratio)
            prev_gray = prev_gray[:, cx:cx+cw]
            curr_gray = curr_gray[:, cx:cx+cw]
            strip_h = max(10, min(int(overlap_px) * 2, prev_h // 2, curr_h - 1, 400))
            needle = prev_gray[prev_h - strip_h: prev_h, :]
            res = cv2.matchTemplate(curr_gray, needle, cv2.TM_CCOEFF_NORMED)
            _, max_val, _, max_loc = cv2.minMaxLoc(res)
            match_y = int(max_loc[1])
            if max_val >= 0.6:
                crop_top = match_y + strip_h - keep_overlap_px - seam_bias_px
            else:
                crop_top = int(overlap_px) - keep_overlap_px - seam_bias_px
        else:
            # Fallback: fixed overlap
            crop_top = int(overlap_px) - keep_overlap_px - seam_bias_px

        crop_top = max(0, min(int(crop_top), curr_h))
        if crop_top >= curr_h - 5:
            # No meaningful new content; stop stitching
            break

        new_part = curr.crop((0, crop_top, curr_w, curr_h))
        # Extend output and paste new part with feathered seam blending
        new_total_h = out.size[1] + new_part.size[1]
        extended = Image.new('RGB', (base_w, new_total_h), (255, 255, 255))
        extended.paste(out, (0, 0))

        # Paste lower portion (excluding seam band) straight below
        seam_h = max(0, int(seam_bias_px // 2))  # small extra stability; keep <= seam_blend area
        seam_band = max(8, 12)  # minimal band to blend even if constants are small
        blend_h = max(seam_band, 12)
        lower = new_part.crop((0, blend_h, base_w, new_part.size[1]))
        extended.paste(lower, (0, out.size[1]))

        # Feather-blend the top band of the new part over the bottom of the old
        try:
            import numpy as _np
            band = new_part.crop((0, 0, base_w, blend_h))
            grad = _np.linspace(0, 255, blend_h, dtype=_np.uint8).reshape(blend_h, 1)
            grad = _np.repeat(grad, base_w, axis=1)
            mask = Image.fromarray(grad, mode='L')
            extended.paste(band, (0, out.size[1] - blend_h), mask)
        except Exception:
            # If numpy is unavailable, fall back to a simple paste that may show a light seam
            extended.paste(new_part.crop((0, 0, base_w, blend_h)), (0, out.size[1] - blend_h))

        out = extended

    out.save(output_path)

    # Restore scroll position back to top by paging up
    try:
        # Ensure focus remains inside content region
        pageups = max(2, len(images))  # roughly mirror down scrolls
        for _ in range(pageups):
            pyautogui.press('pageup')
            time.sleep(0.05)
    except Exception:
        pass


if __name__ == "__main__":
    import argparse

    # Simplified CLI: only this command form is supported
    # python fullpage_screenshot_v2.py OUT.png --use-pagedown --window-title-contains "Chrome" --activate --maximize
    parser = argparse.ArgumentParser(description="Full-page screenshot v2 (scroll+stitch current window)")
    parser.add_argument("out", help="Output PNG path")
    parser.add_argument("--window-title-contains", dest="window_title_contains", default="Chrome", help="Pick window whose title contains this text (default: 'Chrome')")
    parser.add_argument("--use-pagedown", dest="use_pagedown", action="store_true", help="Use PageDown instead of mouse wheel")
    parser.add_argument("--activate", dest="activate", action="store_true", help="Bring target window to front before capture")
    parser.add_argument("--maximize", dest="maximize", action="store_true", help="Maximize window before capture")
    parser.add_argument("--content-top-offset", dest="content_top_offset", type=int, default=0, help="Pixels to crop from window top (0 = no crop)")
    parser.add_argument("--content-bottom-offset", dest="content_bottom_offset", type=int, default=0, help="Pixels to crop from window bottom (0 = no crop)")
    parser.add_argument("--seam-bias", dest="seam_bias", type=int, default=16, help="Extra pixels kept above the seam from the next frame")
    parser.add_argument("--align-center", dest="align_center", action="store_true", help="Match using the center strip only (more stable)")
    # Defaults so the one command works without extra flags if desired
    parser.set_defaults(activate=True, maximize=True)
    args = parser.parse_args()

    save_fullpage_screenshot_v2(
        args.out,
        window_title_contains=args.window_title_contains,
        activate=args.activate,
        maximize=args.maximize,
        # Expose offsets via CLI; defaults to no cropping
        delay_seconds=0.5,
        content_top_offset=args.content_top_offset,
        content_bottom_offset=args.content_bottom_offset,
        overlap_px=160,
        scroll_px=600,
        scroll_pause=1.0,
        max_shots=30,
        use_pagedown=args.use_pagedown,
    )
    print(f"Saved stitched full-page screenshot to {args.out}")