#!/usr/bin/env python3
import argparse
import json
import os
import subprocess
import sys
import time


def run(args):
	return subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)


def image_size(path):
	try:
		from PIL import Image
		with Image.open(path) as img:
			return img.size
	except Exception:
		pass
	result = run(["magick", "identify", "-format", "%w %h", path])
	if result.returncode == 0:
		parts = result.stdout.strip().split()
		if len(parts) == 2:
			return int(parts[0]), int(parts[1])
	return None, None


def wait_for_ready(path, timeout):
	if not path:
		return True
	deadline = time.time() + timeout
	while time.time() < deadline:
		if os.path.exists(path):
			return True
		time.sleep(0.05)
	return False


def find_window(title):
	if not title:
		return None
	result = run(["xdotool", "search", "--name", title])
	if result.returncode != 0:
		return None
	for line in result.stdout.splitlines():
		line = line.strip()
		if line:
			return line
	return None


def capture(window_id, output, root_fallback):
	if window_id:
		result = run(["import", "-window", window_id, output])
		if result.returncode == 0:
			return True, None
		if not root_fallback:
			return False, result.stderr.strip() or "import failed"
	if root_fallback:
		result = run(["import", "-window", "root", output])
		if result.returncode == 0:
			return True, None
		return False, result.stderr.strip() or "root import failed"
	return False, "window not found"


def main():
	parser = argparse.ArgumentParser()
	parser.add_argument("--title", required=True)
	parser.add_argument("--output", required=True)
	parser.add_argument("--ready-file")
	parser.add_argument("--timeout", type=float, default=15.0)
	parser.add_argument("--delay", type=float, default=0.0)
	parser.add_argument("--root-fallback", action="store_true")
	parser.add_argument("--json", dest="json_path")
	args = parser.parse_args()

	os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
	if args.json_path:
		os.makedirs(os.path.dirname(os.path.abspath(args.json_path)), exist_ok=True)

	result = {
		"ok": False,
		"output": args.output,
		"title": args.title,
		"window_id": None,
		"width": None,
		"height": None,
		"error": None,
	}

	if not wait_for_ready(args.ready_file, args.timeout):
		result["error"] = "ready timeout"
	else:
		if args.delay > 0:
			time.sleep(args.delay)
		window_id = find_window(args.title)
		result["window_id"] = window_id
		ok, err = capture(window_id, args.output, args.root_fallback)
		result["ok"] = ok
		result["error"] = err
		if ok:
			result["width"], result["height"] = image_size(args.output)

	data = json.dumps(result, sort_keys=True)
	if args.json_path:
		with open(args.json_path, "w", encoding="utf-8") as f:
			f.write(data + "\n")
	else:
		print(data)
	return 0 if result["ok"] else 1


if __name__ == "__main__":
	sys.exit(main())