import mss
import cv2
import numpy as np
import time
from typing import List, Tuple

# ===== User config =====
monitor_area = {"top": 120, "left": 330, "width": 1900, "height": 1263}
# monitor_area = {"top": 121, "left": 27, "width": 672-27, "height": 549-121}

# Einträge dürfen Pixel (x0,y0,x1,y1) oder normiert [0..1] sein (identisch zur Env)
ui_exclude_rects: List[Tuple[float, float, float, float]] = [
    # (0.0, 0.0, 1.0, 0.08),  # Beispiel: oberer HUD-Streifen (normiert)
]
# Feste Referenz (muss zur Env passen!)
REF_SIZE = (1900, 1263)

# Anzeige-Skalierung (nur fürs Fenster)
WINDOW_SCALE = 0.8
MODE_OVERLAY, MODE_TEXT_ONLY = 0, 1
mode = MODE_OVERLAY

# ===== Implementation =====
# Wichtig: Die Detection/Parameter kommen aus der Env, um Drift zu vermeiden.
from flower_game_env import FlowerGameEnv  # Datei muss als flower_game_env.py vorliegen

env = FlowerGameEnv(
    monitor_area,
    ui_exclude_rects=ui_exclude_rects,
    ref_size=REF_SIZE,
)


def _draw_overlay(frame_bgr, det, fps: float):
    """Zeichnet exakt die Größen, die die Env auch nutzt (achsenweise Skalierung)."""
    h, w = det["frame_hw"]

    # Rechteck-Halbachsen in Pixel aus achsenweisen nd-Schwellen
    eat_half_w_px  = int(round((env.eat_x_nd       * w) * 0.5))
    eat_half_h_px  = int(round((env.eat_y_nd       * h) * 0.5))
    coll_half_w_px = int(round((env.collision_x_nd * w) * 0.5))
    coll_half_h_px = int(round((env.collision_y_nd * h) * 0.5))

    tx, ty = det["turtle_xy"]
    fx, fy = det["flower_xy"]
    bombs  = det["bombs_xy"]

    # Nearest Bomb (aus state_norm, schon normiert)
    n_tx, n_ty, n_fx, n_fy, nbx, nby = det["state_norm"]
    nb_px = (int(round(nbx * w)), int(round(nby * h))) if (len(bombs) > 0) else None

    # Flower
    if fx is not None:
        cv2.circle(frame_bgr, (fx, fy), 8, (0, 255, 255), 2)
        cv2.putText(frame_bgr, "Flower", (fx+10, fy-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,255), 2)

    # Bomben (alle grau, nächste rot)
    for (bx, by, _a) in bombs:
        cv2.circle(frame_bgr, (bx, by), 10, (60, 60, 60), 2)
    if nb_px is not None:
        cv2.circle(frame_bgr, nb_px, 12, (0, 0, 255), 3)

    # Turtle + Zonenrechtecke
    if tx is not None:
        cv2.circle(frame_bgr, (tx, ty), 8, (0, 200, 0), 2)
        cv2.putText(frame_bgr, "Turtle", (tx+10, ty-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,200,0), 2)
        cv2.rectangle(frame_bgr, (tx - eat_half_w_px,  ty - eat_half_h_px),
                      (tx + eat_half_w_px,  ty + eat_half_h_px),  (0,255,0), 1)
        cv2.rectangle(frame_bgr, (tx - coll_half_w_px, ty - coll_half_h_px),
                      (tx + coll_half_w_px, ty + coll_half_h_px), (0,0,255), 1)

    # HUD
    min_area_px = int(round(env.bomb_min_area_frac * h * w))
    hud_lines = [
        f"FPS: {fps:.1f}",
        f"state = [tx={n_tx:.3f}, ty={n_ty:.3f}, fx={n_fx:.3f}, fy={n_fy:.3f}, bx={nbx:.3f}, by={nby:.3f}]",
        f"eat_x_nd={env.eat_x_nd:.5f}  eat_y_nd={env.eat_y_nd:.5f}  coll_x_nd={env.collision_x_nd:.5f}  coll_y_nd={env.collision_y_nd:.5f}",
        f"bomb_min_area_frac={env.bomb_min_area_frac:.7f} (-> min_area_px≈{min_area_px})  ref_size={env.w_ref}x{env.h_ref}",
    ]
    y = 28
    for line in hud_lines:
        cv2.putText(frame_bgr, line, (16, y), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255,255,255), 2)
        y += 24

    return frame_bgr


def main():
    global mode
    sct = mss.mss()
    prev = time.time()
    fps = 0.0

    while True:
        raw = np.array(sct.grab(monitor_area))  # BGRA
        frame = cv2.cvtColor(raw, cv2.COLOR_BGRA2BGR)

        # Detection mit exakt derselben Logik wie in der Env
        det = env._detect_entities(frame)

        now = time.time()
        dt = now - prev; prev = now
        if dt > 0:
            fps = 1.0 / dt

        if mode == MODE_OVERLAY:
            out = _draw_overlay(frame.copy(), det, fps)
        else:
            h, w = det["frame_hw"]
            out = np.zeros((h, w, 3), dtype=np.uint8)
            out = _draw_overlay(out, det, fps)

        if WINDOW_SCALE != 1.0:
            out = cv2.resize(out, (int(out.shape[1] * WINDOW_SCALE), int(out.shape[0] * WINDOW_SCALE)))
        cv2.imshow("Env-State Viewer (axes-normalized)", out)

        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'): break
        elif key == ord('0'): mode = MODE_OVERLAY
        elif key == ord('1'): mode = MODE_TEXT_ONLY

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()