"""
TelecomKz Long Screenshot & Multi-screen Capture Script
Captures continuous scroll screenshots via ADB and stitches them cleanly.
"""

import sys
import subprocess
import time
from pathlib import Path
from PIL import Image
import io

ADB_PATHS = [
    r"C:\Users\user\AppData\Local\Android\Sdk\platform-tools\adb.exe",
    r"C:\Program Files\Netease\MuMuPlayer\nx_main\adb.exe",
    "adb"
]

def find_adb():
    for p in ADB_PATHS:
        try:
            res = subprocess.run([p, "version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            if res.returncode == 0:
                return p
        except Exception:
            continue
    return "adb"

ADB_BIN = find_adb()

def get_devices():
    res = subprocess.run([ADB_BIN, "devices"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    lines = res.stdout.strip().split("\n")[1:]
    devices = []
    for line in lines:
        parts = line.strip().split("\t")
        if len(parts) >= 2 and parts[1] == "device":
            devices.append(parts[0])
    return devices

def capture_single_image(device_id=None):
    cmd = [ADB_BIN]
    if device_id:
        cmd.extend(["-s", device_id])
    cmd.extend(["exec-out", "screencap", "-p"])
    res = subprocess.run(cmd, stdout=subprocess.PIPE)
    if res.returncode == 0 and len(res.stdout) > 1000:
        return Image.open(io.BytesIO(res.stdout))
    return None

def capture_long_screenshot(output_path: str, scroll_steps: int = 3, device_id: str = None):
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    images = []
    
    print("Capturing top part of screen...")
    img1 = capture_single_image(device_id)
    if not img1:
        print("Failed to capture screen")
        return False
    images.append(img1)
    
    w, h = img1.size
    # Status bar top ~ 120px, bottom nav bar ~ 200px
    top_crop = 250 # Below top bar
    bottom_crop = 200 # Above bottom bar
    effective_h = h - top_crop - bottom_crop
    
    for i in range(scroll_steps):
        print(f"Scrolling step {i+1}/{scroll_steps}...")
        # Swipe up (scroll down)
        subprocess.run([ADB_BIN, "shell", "input", "swipe", "540", "1600", "540", "750", "400"])
        time.sleep(1.2)
        
        img = capture_single_image(device_id)
        if img:
            images.append(img)
            
    # If 1 image only
    if len(images) == 1:
        images[0].save(output_path)
        print(f"Single image saved to: {output_path}")
        return True
        
    # Stitch images vertically
    # First image: full or crop
    total_stitched_height = h + (len(images) - 1) * (effective_h - 100)
    stitched = Image.new("RGBA", (w, total_stitched_height), (10, 15, 29, 255))
    
    # Paste 1st image
    stitched.paste(images[0], (0, 0))
    current_y = h - bottom_crop
    
    for i in range(1, len(images)):
        # Crop the middle content of the scrolled image
        crop_box = (0, top_crop, w, h - bottom_crop)
        cropped_slice = images[i].crop(crop_box)
        stitched.paste(cropped_slice, (0, current_y))
        current_y += cropped_slice.size[1] - 100
        
    stitched = stitched.crop((0, 0, w, current_y + bottom_crop))
    stitched.save(output_path, "PNG")
    print(f"Long stitched screenshot saved to: {output_path} (Size: {stitched.size[0]}x{stitched.size[1]})")
    return True

def main():
    devices = get_devices()
    if not devices:
        print("No Android device connected.")
        sys.exit(1)
    dev = devices[0]
    
    # Capture single clean screenshot
    print("Capturing standard high-res screenshot...")
    img = capture_single_image(dev)
    if img:
        img.save("public/assets/screens/telecomkz_main.png")
        print(f"Main screen saved: {img.size[0]}x{img.size[1]}")
        
    # Also capture stitched long screenshot
    print("Capturing stitched long screenshot...")
    capture_long_screenshot("public/assets/screens/telecomkz_main_long.png", scroll_steps=2, device_id=dev)

if __name__ == "__main__":
    main()