DirtyLeague_Bot/calibrate_tool.py

"""
Calibration and ROI helper tool.
Used to:
- Capture current game screen and save snapshots.
- Crop template snippets (buttons, banners) for assets/
- Test template matching against the active game screen with confidence score.
"""

import argparse
import json
import os
import sys
import time
import cv2
import numpy as np
from window_utils import WindowManager

CONFIG_PATH = os.path.join(os.path.dirname(__file__), "config.json")


def load_config():
    if os.path.exists(CONFIG_PATH):
        with open(CONFIG_PATH, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}


def capture_snapshot(output_path: str = "assets/current_screen.png"):
    config = load_config()
    target_title = config.get("target_window", {}).get("title", "DirtyLeague")
    wm = WindowManager(window_title=target_title)

    hwnd = wm.find_window()
    if not hwnd:
        print(f"[ERROR] Window '{target_title}' not found!")
        return None

    frame = wm.capture_frame(client_only=True)
    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
    cv2.imwrite(output_path, frame)
    h, w = frame.shape[:2]
    print(f"[OK] Screenshot saved to: {output_path} (Resolution: {w}x{h})")
    return output_path


def crop_roi(source_path: str, output_name: str, x: int, y: int, w: int, h: int):
    if not os.path.exists(source_path):
        print(f"[ERROR] Source file not found: {source_path}")
        return

    img = cv2.imread(source_path)
    if img is None:
        print(f"[ERROR] Could not read image: {source_path}")
        return

    roi = img[y:y+h, x:x+w]
    if roi.size == 0:
        print(f"[ERROR] Crop dimensions resulted in empty image! (x={x}, y={y}, w={w}, h={h})")
        return

    out_path = os.path.join("assets", output_name if output_name.endswith(".png") else f"{output_name}.png")
    cv2.imwrite(out_path, roi)
    print(f"[OK] Cropped ROI saved to: {out_path} ({w}x{h})")


def test_match(template_name: str, threshold: float = 0.8):
    config = load_config()
    target_title = config.get("target_window", {}).get("title", "DirtyLeague")
    wm = WindowManager(window_title=target_title)

    template_path = os.path.join("assets", template_name if template_name.endswith(".png") else f"{template_name}.png")
    if not os.path.exists(template_path):
        print(f"[ERROR] Template not found: {template_path}")
        return

    tpl = cv2.imread(template_path)
    if tpl is None:
        print(f"[ERROR] Failed to load template: {template_path}")
        return

    frame = wm.capture_frame(client_only=True)
    res = cv2.matchTemplate(frame, tpl, cv2.TM_CCOEFF_NORMED)
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)

    th, tw = tpl.shape[:2]
    print(f"--- Template Match: {template_name} ---")
    print(f"Max confidence: {max_val:.4f} (Threshold: {threshold})")
    print(f"Best match client location: ({max_loc[0]}, {max_loc[1]}) -> Center: ({max_loc[0] + tw // 2}, {max_loc[1] + th // 2})")

    if max_val >= threshold:
        print("[RESULT] MATCH FOUND! Target is visible.")
    else:
        print("[RESULT] NOT FOUND (below threshold).")


def check_all_assets():
    config = load_config()
    target_title = config.get("target_window", {}).get("title", "DirtyLeague")
    default_conf = config.get("bot", {}).get("confidence_threshold", 0.8)
    wm = WindowManager(window_title=target_title)

    assets_dir = "assets"
    if not os.path.exists(assets_dir):
        print("[ERROR] assets/ directory does not exist.")
        return

    files = [f for f in os.listdir(assets_dir) if f.endswith(".png") and not f.startswith("current_screen")]
    if not files:
        print("[INFO] No template images found in assets/ yet.")
        return

    frame = wm.capture_frame(client_only=True)
    print(f"Testing {len(files)} asset templates against current game screen...")
    for f in sorted(files):
        tpl = cv2.imread(os.path.join(assets_dir, f))
        if tpl is None:
            continue
        res = cv2.matchTemplate(frame, tpl, cv2.TM_CCOEFF_NORMED)
        _, max_val, _, max_loc = cv2.minMaxLoc(res)
        status = "MATCH" if max_val >= default_conf else "NO MATCH"
        print(f"  [{status}] {f:<25} confidence: {max_val:.3f} loc: {max_loc}")


def main():
    parser = argparse.ArgumentParser(description="Bot UI Calibration & Asset Tool")
    subparsers = parser.add_subparsers(dest="command")

    # snapshot command
    snap_p = subparsers.add_parser("snapshot", help="Capture current game screen to file")
    snap_p.add_argument("-o", "--output", default="assets/current_screen.png", help="Output PNG path")

    # crop command
    crop_p = subparsers.add_parser("crop", help="Crop ROI template from an image")
    crop_p.add_argument("-s", "--source", default="assets/current_screen.png", help="Source screenshot")
    crop_p.add_argument("-n", "--name", required=True, help="Output template name (e.g. btn_play.png)")
    crop_p.add_argument("--roi", nargs=4, type=int, required=True, metavar=("X", "Y", "W", "H"), help="Crop region X Y W H")

    # test command
    test_p = subparsers.add_parser("test", help="Test match of a single template on live screen")
    test_p.add_argument("-n", "--name", required=True, help="Template name in assets/")
    test_p.add_argument("-t", "--threshold", type=float, default=0.8, help="Confidence threshold")

    # scan command
    subparsers.add_parser("scan", help="Scan all templates in assets/ against current screen")

    args = parser.parse_args()
    if args.command == "snapshot":
        capture_snapshot(args.output)
    elif args.command == "crop":
        x, y, w, h = args.roi
        crop_roi(args.source, args.name, x, y, w, h)
    elif args.command == "test":
        test_match(args.name, args.threshold)
    elif args.command == "scan":
        check_all_assets()
    else:
        parser.print_help()


if __name__ == "__main__":
    main()