from __future__ import annotations

"""
Extract a frame JSON (hand cards + elixir) from a BlueStacks screenshot, using:
- local OCR for elixir (calibrated ROI)
- local template matching for hand cards (deck icons from clash_royale_pack_for_mac/icons)

This is inspired by the schema in:
  clash_royale_pack_for_mac/outputs/frame_with_templates.json

Usage:
  cd toxic_royale_env
  python3 scripts/bluestacks_frame_extractor.py --image outputs/bluestacks/pilot_current.png

Debug:
  BS_DEBUG_FRAME=1 python3 scripts/bluestacks_frame_extractor.py --image ...
  -> writes outputs/bluestacks/debug_hand_slot_*.png
"""

import argparse
import json
import os
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from PIL import Image


def _load_json(path: Path) -> dict:
    return json.loads(path.read_text(encoding="utf-8"))


def _load_elixir_reader(root: Path):
    # Local import without package requirements
    import importlib.util
    import sys

    mod_path = root / "scripts" / "bluestacks_elixir_ocr.py"
    spec = importlib.util.spec_from_file_location("bluestacks_elixir_ocr", mod_path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Could not load module from {mod_path}")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    fn = getattr(module, "read_elixir", None)
    if fn is None:
        raise RuntimeError("Missing read_elixir() in bluestacks_elixir_ocr.py")
    return fn


def _bluestacks_bounds() -> tuple[int, int, int, int] | None:
    """
    Must match scripts/bluestacks_screenshot.sh bounds logic.
    """
    import subprocess

    script = r'''
      tell application "System Events"
        try
          set bluestacksProc to application process "BlueStacks"
          set allWins to every window of bluestacksProc
          repeat with w in allWins
            set t to title of w
            set n to name of w
            set sr to subrole of w
            set s to size of w
            set p to position of w
            if (t is "BlueStacks Air" or n is "BlueStacks Air") and sr is "AXDialog" then
              set x to item 1 of p
              set y to item 2 of p
              set w_ to item 1 of s
              set h to item 2 of s
              if h > 200 then
                return "" & x & "," & y & "," & w_ & "," & h
              end if
            end if
          end repeat
          return ""
        on error
          return ""
        end try
      end tell
    '''
    try:
        out = subprocess.check_output(["osascript", "-e", script], text=True).strip()
        if not out or out.count(",") < 3:
            return None
        x_s, y_s, w_s, h_s = out.split(",", 3)
        return int(float(x_s)), int(float(y_s)), int(float(w_s)), int(float(h_s))
    except Exception:
        return None


@dataclass(frozen=True)
class Match:
    card: str
    conf: float


ICON_MATCH_SIZE = (96, 96)


def _prep_gray(im: Image.Image, *, top_drop_frac: float) -> Image.Image:
    """
    Make matching robust to borders/UI by:
    - dropping some of the top region
    - center-cropping
    - converting to grayscale
    """
    w, h = im.size
    y0 = int(round(h * top_drop_frac))
    im0 = im.crop((0, y0, w, h))

    w2, h2 = im0.size
    cx0 = int(round(w2 * 0.10))
    cy0 = int(round(h2 * 0.10))
    cx1 = int(round(w2 * 0.90))
    cy1 = int(round(h2 * 0.95))
    return im0.crop((cx0, cy0, cx1, cy1)).convert("L")


def _to_norm_vec(im: Image.Image, *, top_drop_frac: float) -> np.ndarray:
    g = _prep_gray(im, top_drop_frac=top_drop_frac).resize(ICON_MATCH_SIZE, Image.Resampling.BICUBIC)
    arr = np.asarray(g, dtype=np.float32).reshape(-1)
    arr -= float(arr.mean())
    n = float(np.linalg.norm(arr) + 1e-6)
    return arr / n


def _load_templates(pack_dir: Path) -> dict[str, np.ndarray]:
    """
    Loads deck icons as grayscale arrays (no evo / no card_templates).
    """
    icon_dir = pack_dir / "icons"
    out: dict[str, np.ndarray] = {}

    # Deck-only matching (reduces false positives a lot).
    want = {
        "knight",
        "goblins",
        "minions",
        "bomber",
        "fireball",
        "goblin-hut",
        "musketeer",
        "mini-pekka",
    }

    def _icon_path_for(key: str) -> Path:
        # most are snake_case
        base = key.replace("-", "_")
        candidates = [
            f"{base}.png",
            # special naming in this pack
            "mini_p_e_k_k_a.png" if key == "mini-pekka" else "",
            "p_e_k_k_a.png" if key == "pekka" else "",
        ]
        for name in candidates:
            if not name:
                continue
            p = icon_dir / name
            if p.exists():
                return p
        return icon_dir / f"{base}.png"

    # Icons are always available for these.
    for key in sorted(want):
        ip = _icon_path_for(key)
        if ip.exists():
            out[key] = _to_norm_vec(Image.open(ip), top_drop_frac=0.0)
        else:
            raise RuntimeError(f"missing_deck_icon::{ip}")
    return out


def _hand_slot_boxes_px(cfg: dict, *, img_w: int, img_h: int) -> list[tuple[int, int, int, int]]:
    """
    Compute 4 hand slot crop boxes in screenshot pixel coords.

    We use calibrated card_slots centers (screen coords) and map to screenshot pixels
    via BlueStacks bounds + Retina scale.
    """
    b = _bluestacks_bounds()
    if b is None:
        raise RuntimeError("cannot_find_bluestacks_bounds")
    ox, oy, w_pt, h_pt = b
    sx = img_w / max(1.0, float(w_pt))
    sy = img_h / max(1.0, float(h_pt))

    centers = []
    for k in ["1", "2", "3", "4"]:
        x_s, y_s = (cfg.get("card_slots") or {}).get(k) or [0, 0]
        cx = int(round((int(x_s) - ox) * sx))
        cy = int(round((int(y_s) - oy) * sy))
        centers.append((cx, cy))

    # Estimate slot width from distances between centers.
    dxs = [abs(centers[i + 1][0] - centers[i][0]) for i in range(3)]
    dx = int(round(float(sorted(dxs)[1]) if len(dxs) >= 3 else (dxs[0] if dxs else 120)))
    # Card art crop: target the top artwork region (avoid the bottom cost bubble).
    # We use a square crop and shift it upward relative to the slot center.
    w_box = int(round(dx * 0.78))
    h_box = w_box
    y_up = int(round(h_box * 0.78))
    y_dn = int(round(h_box * 0.22))
    y_shift_up = int(round(dx * 0.22))

    boxes = []
    for cx, cy in centers:
        cy2 = cy - y_shift_up
        x0 = max(0, cx - w_box // 2)
        x1 = min(img_w, cx + w_box // 2)
        y0 = max(0, cy2 - y_up)
        y1 = min(img_h, cy2 + y_dn)
        boxes.append((x0, y0, x1, y1))
    return boxes


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--image", type=str, required=True)
    args = ap.parse_args()

    root = Path(__file__).resolve().parents[1]
    pack_dir = root / "clash_royale_pack_for_mac"
    cfg = _load_json(root / "config" / "bluestacks_gameplay.local.json")

    img_path = Path(args.image)
    im = Image.open(img_path).convert("RGB")

    read_elixir = _load_elixir_reader(root)
    elixir, el_dbg = read_elixir(img_path)

    # templates
    templates = _load_templates(pack_dir)

    boxes = _hand_slot_boxes_px(cfg, img_w=im.size[0], img_h=im.size[1])

    debug = os.environ.get("BS_DEBUG_FRAME", "0") == "1"
    if debug:
        out_dir = root / "outputs" / "bluestacks"
        out_dir.mkdir(parents=True, exist_ok=True)

    hand = []
    for idx, box in enumerate(boxes):
        crop = im.crop(box)
        search_vec = _to_norm_vec(crop, top_drop_frac=0.18)

        best = Match(card="unknown", conf=-1.0)
        for card, tpl_vec in templates.items():
            s = float(np.dot(search_vec, tpl_vec))
            if s > best.conf:
                best = Match(card=card, conf=s)

        # Map template key -> dataset key/cost
        # Our pack uses kebab-case stems; dataset.json has names but we can map for known 8.
        cost_map = {
            "knight": 3,
            "goblins": 2,
            "minions": 3,
            "bomber": 2,
            "fireball": 4,
            "goblin-hut": 4,
            "musketeer": 4,
            "mini-pekka": 4,
            "archers": 3,
            "arrows": 3,
            "giant": 5,
        }
        cost = cost_map.get(best.card)
        playable = (elixir is not None) and (cost is not None) and (float(cost) <= float(elixir))

        hand.append(
            {
                "slot": idx,
                "card": best.card,
                "cost": cost,
                "is_playable": bool(playable),
                "cooldown_ms": 0,
                "bbox_x1": int(box[0]),
                "bbox_y1": int(box[1]),
                "bbox_x2": int(box[2]),
                "bbox_y2": int(box[3]),
                "conf": float(best.conf),
            }
        )

        if debug:
            out_dir = root / "outputs" / "bluestacks"
            crop.save(out_dir / f"debug_hand_slot_{idx}_{best.card}.png")

    frame = {
        "meta": {"source": {"device": "mac", "emulator": "bluestacks", "game": "clash_royale"}},
        "ui": {"screen": "battle"},
        "player": {
            "elixir": elixir,
            "hand": hand,
            "debug_elixir": el_dbg,
        },
    }

    print(json.dumps(frame, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())