import argparse
import json
import signal
import subprocess
import threading
import time
import webbrowser
from datetime import datetime, timezone
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from io import BytesIO
from pathlib import Path

import numpy as np
import onnxruntime as ort
from PIL import Image, ImageDraw

import Quartz

from ffmpeg_realtime_detect import (
    CLASS_NAMES,
    MODEL_PATH,
    detect,
    detection_event,
    draw_detections,
    image_to_jpeg_bytes,
    no_detection_event,
)


class StreamState:
    def __init__(self):
        self.condition = threading.Condition()
        self.jpeg = None
        self.event = None
        self.sequence = 0

    def update(self, jpeg, event):
        with self.condition:
            self.jpeg = jpeg
            self.event = event
            self.sequence += 1
            self.condition.notify_all()

    def snapshot(self):
        with self.condition:
            return self.jpeg, self.event, self.sequence


def parse_args():
    parser = argparse.ArgumentParser(
        description="Capture only the WeChat window, run ONNX, and expose a live stream."
    )
    parser.add_argument("--model", default=str(MODEL_PATH), help="ONNX model path")
    parser.add_argument("--fps", type=float, default=2.0, help="capture FPS")
    parser.add_argument(
        "--infer-interval",
        type=float,
        default=3.0,
        help="seconds between ONNX inference runs; preview frames reuse the latest detections",
    )
    parser.add_argument(
        "--log-every",
        type=int,
        default=60,
        help="log every N preview frames in addition to inference frames; 0 disables periodic logs",
    )
    parser.add_argument("--confidence", type=float, default=0.05, help="minimum score")
    parser.add_argument("--target-class", type=int, default=0, help="target class id")
    parser.add_argument("--host", default="127.0.0.1", help="HTTP host")
    parser.add_argument("--port", type=int, default=8765, help="HTTP port")
    parser.add_argument("--max-frames", type=int, default=0, help="stop after N frames; 0 runs forever")
    parser.add_argument("--open-browser", action="store_true", help="open the live preview in browser")
    parser.add_argument("--open-ffplay", action="store_true", help="open ffplay for the MJPEG stream")
    return parser.parse_args()


def utc_now():
    return datetime.now(timezone.utc).isoformat()


def emit(event):
    print(json.dumps(event, ensure_ascii=False, separators=(",", ":")), flush=True)


def find_wechat_window():
    options = Quartz.kCGWindowListOptionOnScreenOnly | Quartz.kCGWindowListExcludeDesktopElements
    windows = Quartz.CGWindowListCopyWindowInfo(options, Quartz.kCGNullWindowID) or []
    for window in windows:
        owner = window.get(Quartz.kCGWindowOwnerName, "") or ""
        title = window.get(Quartz.kCGWindowName, "") or ""
        layer = window.get(Quartz.kCGWindowLayer, -1)
        if layer != 0:
            continue
        if owner == "微信" or owner == "WeChat" or "wechat" in owner.lower() or title == "微信":
            bounds = window.get(Quartz.kCGWindowBounds, {}) or {}
            return {
                "id": int(window.get(Quartz.kCGWindowNumber)),
                "owner": owner,
                "title": title,
                "x": float(bounds.get("X", 0)),
                "y": float(bounds.get("Y", 0)),
                "width": float(bounds.get("Width", 0)),
                "height": float(bounds.get("Height", 0)),
            }
    return None


def capture_window_image(window_id):
    image_ref = Quartz.CGWindowListCreateImage(
        Quartz.CGRectNull,
        Quartz.kCGWindowListOptionIncludingWindow,
        window_id,
        Quartz.kCGWindowImageBoundsIgnoreFraming,
    )
    if image_ref is None:
        return None

    width = Quartz.CGImageGetWidth(image_ref)
    height = Quartz.CGImageGetHeight(image_ref)
    bytes_per_row = Quartz.CGImageGetBytesPerRow(image_ref)
    provider = Quartz.CGImageGetDataProvider(image_ref)
    data = Quartz.CGDataProviderCopyData(provider)
    array = np.frombuffer(data, dtype=np.uint8)
    array = array.reshape((height, bytes_per_row))[:, : width * 4]
    bgra = array.reshape((height, width, 4))
    rgba = bgra[:, :, [2, 1, 0, 3]]
    return Image.fromarray(rgba, "RGBA").convert("RGB")


def screen_event(event, window):
    event = dict(event)
    event["window"] = window
    frame_width, frame_height = event.get("frame_size") or [window["width"], window["height"]]
    scale_x = frame_width / max(window["width"], 1)
    scale_y = frame_height / max(window["height"], 1)
    if event.get("bbox_screen"):
        x1, y1, x2, y2 = event["bbox_screen"]
        event["bbox_image"] = [x1, y1, x2, y2]
        wx1, wy1, wx2, wy2 = x1 / scale_x, y1 / scale_y, x2 / scale_x, y2 / scale_y
        event["bbox_window"] = [round(wx1, 2), round(wy1, 2), round(wx2, 2), round(wy2, 2)]
        event["bbox_screen"] = [
            round(window["x"] + wx1, 2),
            round(window["y"] + wy1, 2),
            round(window["x"] + wx2, 2),
            round(window["y"] + wy2, 2),
        ]
    if event.get("center_screen"):
        cx, cy = event["center_screen"]
        event["center_image"] = [cx, cy]
        wcx, wcy = cx / scale_x, cy / scale_y
        event["center_window"] = [round(wcx, 2), round(wcy, 2)]
        event["center_screen"] = [round(window["x"] + wcx, 2), round(window["y"] + wcy, 2)]
    return event


def event_from_detections(detections, frame_index, frame_size, class_name, target_class):
    if len(detections) > 0:
        best = detections[np.argmax(detections[:, 4])]
        return detection_event(best, frame_index, frame_size, class_name)
    return no_detection_event(frame_index, frame_size, target_class, class_name)


def draw_overlay(image, event):
    draw = ImageDraw.Draw(image)
    text = f"frame={event.get('frame_index', '-')} {event.get('timestamp', '')}"
    draw.rectangle((12, 12, 720, 56), fill=(0, 0, 0))
    draw.text((22, 24), text, fill=(80, 255, 120))
    return image


def html_page():
    return """<!doctype html>
<html lang="zh-CN">
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1" />
  <title>WeChat Window Live</title>
  <style>
    body { margin: 0; background: #09090b; color: #f4f4f5; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }
    .wrap { display: grid; grid-template-columns: minmax(0, 1fr) 360px; gap: 16px; padding: 16px; }
    .stage { background: #18181b; border: 1px solid #27272a; border-radius: 14px; overflow: hidden; }
    img { display: block; width: 100%; height: auto; }
    aside { background: #18181b; border: 1px solid #27272a; border-radius: 14px; padding: 16px; }
    h1 { margin: 0 0 12px; font-size: 18px; }
    .item { padding: 10px 0; border-top: 1px solid #27272a; }
    .label { color: #a1a1aa; font-size: 12px; margin-bottom: 4px; }
    .value { font-family: ui-monospace, SFMono-Regular, Menlo, monospace; word-break: break-all; }
    .ok { color: #86efac; }
    .miss { color: #fca5a5; }
    @media (max-width: 900px) { .wrap { grid-template-columns: 1fr; } }
  </style>
</head>
<body>
  <div class="wrap">
    <main class="stage"><img src="/stream.mjpg" alt="WeChat window live stream" /></main>
    <aside>
      <h1>WeChat Window Live</h1>
      <div class="item"><div class="label">Stream</div><div class="value">/stream.mjpg</div></div>
      <div class="item"><div class="label">Sequence</div><div id="sequence" class="value">-</div></div>
      <div class="item"><div class="label">Status</div><div id="status" class="value">-</div></div>
      <div class="item"><div class="label">Confidence</div><div id="confidence" class="value">-</div></div>
      <div class="item"><div class="label">Center Screen</div><div id="center" class="value">-</div></div>
      <div class="item"><div class="label">BBox Screen</div><div id="bbox" class="value">-</div></div>
      <div class="item"><div class="label">Window</div><div id="window" class="value">-</div></div>
    </aside>
  </div>
  <script>
    async function refresh() {
      try {
        const res = await fetch('/latest.json?t=' + Date.now(), { cache: 'no-store' });
        const data = await res.json();
        const event = data.event || {};
        const hit = event.type === 'detection';
        document.getElementById('sequence').textContent = data.sequence ?? '-';
        const status = document.getElementById('status');
        status.textContent = event.type ?? 'waiting';
        status.className = 'value ' + (hit ? 'ok' : 'miss');
        document.getElementById('confidence').textContent = event.confidence ?? '-';
        document.getElementById('center').textContent = JSON.stringify(event.center_screen ?? '-');
        document.getElementById('bbox').textContent = JSON.stringify(event.bbox_screen ?? '-');
        document.getElementById('window').textContent = JSON.stringify(event.window ?? '-');
      } catch (e) {}
    }
    refresh();
    setInterval(refresh, 300);
  </script>
</body>
</html>
"""


def make_handler(state):
    class Handler(BaseHTTPRequestHandler):
        def log_message(self, format, *args):
            return

        def do_HEAD(self):
            if self.path == "/" or self.path.startswith("/index.html"):
                self.send_response(200)
                self.send_header("Content-Type", "text/html; charset=utf-8")
                self.send_header("Cache-Control", "no-store")
                self.end_headers()
                return
            if self.path.startswith("/latest.json"):
                self.send_response(200)
                self.send_header("Content-Type", "application/json; charset=utf-8")
                self.send_header("Cache-Control", "no-store")
                self.end_headers()
                return
            if self.path.startswith("/stream.mjpg"):
                self.send_response(200)
                self.send_header("Cache-Control", "no-cache, private")
                self.send_header("Pragma", "no-cache")
                self.send_header("Content-Type", "multipart/x-mixed-replace; boundary=frame")
                self.end_headers()
                return
            self.send_error(404)

        def do_GET(self):
            if self.path == "/" or self.path.startswith("/index.html"):
                data = html_page().encode("utf-8")
                self.send_response(200)
                self.send_header("Content-Type", "text/html; charset=utf-8")
                self.send_header("Cache-Control", "no-store")
                self.send_header("Content-Length", str(len(data)))
                self.end_headers()
                self.wfile.write(data)
                return
            if self.path.startswith("/latest.json"):
                _, event, sequence = state.snapshot()
                payload = {"sequence": sequence, "updated_at": utc_now(), "event": event or {"type": "waiting"}}
                data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
                self.send_response(200)
                self.send_header("Content-Type", "application/json; charset=utf-8")
                self.send_header("Cache-Control", "no-store")
                self.send_header("Content-Length", str(len(data)))
                self.end_headers()
                self.wfile.write(data)
                return
            if self.path.startswith("/stream.mjpg"):
                self.send_response(200)
                self.send_header("Cache-Control", "no-cache, private")
                self.send_header("Pragma", "no-cache")
                self.send_header("Content-Type", "multipart/x-mixed-replace; boundary=frame")
                self.end_headers()
                last_sequence = -1
                while True:
                    with state.condition:
                        state.condition.wait_for(lambda: state.sequence != last_sequence, timeout=5)
                        jpeg, _, last_sequence = state.jpeg, state.event, state.sequence
                    if not jpeg:
                        continue
                    try:
                        self.wfile.write(b"--frame\r\n")
                        self.wfile.write(b"Content-Type: image/jpeg\r\n")
                        self.wfile.write(f"Content-Length: {len(jpeg)}\r\n\r\n".encode("ascii"))
                        self.wfile.write(jpeg)
                        self.wfile.write(b"\r\n")
                    except (BrokenPipeError, ConnectionResetError):
                        break
                return
            self.send_error(404)

    return Handler


def start_server(host, port, state):
    server = ThreadingHTTPServer((host, port), make_handler(state))
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    return server


def open_ffplay(url):
    return subprocess.Popen(
        ["ffplay", "-hide_banner", "-loglevel", "error", "-fflags", "nobuffer", "-flags", "low_delay", "-framedrop", url],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )


def run(args):
    state = StreamState()
    server = start_server(args.host, args.port, state)
    url = f"http://{args.host}:{args.port}/"
    stream_url = f"http://{args.host}:{args.port}/stream.mjpg"
    emit({"type": "server_started", "url": url, "stream_url": stream_url})
    if args.open_browser:
        webbrowser.open(url)
    ffplay_process = open_ffplay(stream_url) if args.open_ffplay else None

    session = ort.InferenceSession(str(Path(args.model)))
    input_meta = session.get_inputs()[0]
    class_name = CLASS_NAMES.get(args.target_class, f"class_{args.target_class}")
    interval = 1 / max(args.fps, 0.1)
    frame_index = 0
    last_infer_at = 0.0
    last_infer_size = None
    last_detections = np.empty((0, 6))
    try:
        while args.max_frames <= 0 or frame_index < args.max_frames:
            start = time.time()
            window = find_wechat_window()
            if not window:
                event = {"type": "wechat_window_not_found", "timestamp": utc_now()}
                state.update(image_to_jpeg_bytes(Image.new("RGB", (960, 540), "black")), event)
                emit(event)
                time.sleep(interval)
                continue

            image = capture_window_image(window["id"])
            if image is None:
                event = {"type": "capture_failed", "timestamp": utc_now(), "window": window}
                emit(event)
                time.sleep(interval)
                continue

            frame_index += 1
            now = time.time()
            should_infer = (
                now - last_infer_at >= args.infer_interval
                or last_infer_size != image.size
                or frame_index == 1
            )
            if should_infer:
                last_detections = detect(session, input_meta, image, args.target_class, args.confidence)
                last_infer_at = now
                last_infer_size = image.size

            event = event_from_detections(
                last_detections,
                frame_index,
                image.size,
                class_name,
                args.target_class,
            )
            event["inferred"] = should_infer
            event["inference_age_ms"] = round((now - last_infer_at) * 1000, 1)
            event = screen_event(event, window)
            annotated = draw_detections(image, last_detections, class_name)
            annotated = draw_overlay(annotated, event)
            state.update(image_to_jpeg_bytes(annotated), event)
            if should_infer or (args.log_every > 0 and frame_index % args.log_every == 0):
                emit(event)

            elapsed = time.time() - start
            time.sleep(max(0, interval - elapsed))
    finally:
        server.shutdown()
        if ffplay_process:
            ffplay_process.terminate()


def main():
    def stop_handler(signum, frame):
        raise KeyboardInterrupt

    signal.signal(signal.SIGTERM, stop_handler)
    signal.signal(signal.SIGINT, stop_handler)
    args = parse_args()
    try:
        run(args)
    except KeyboardInterrupt:
        emit({"type": "stopped", "timestamp": utc_now()})


if __name__ == "__main__":
    main()
