import argparse
import base64
import json
import signal
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
from urllib import request as urlrequest

import numpy as np
from PIL import Image, ImageDraw
import pyautogui
import pyperclip
import tomllib

from wechat_window_live import capture_window_image, find_wechat_window, image_to_jpeg_bytes


DEFAULT_REGIONS = Path.home() / "Library/Application Support/com.tauri.dev/data/regions/wechat.json"
DEFAULT_AGENT_CONFIG = Path(__file__).resolve().parent.parent / "agent" / "config.toml"


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


class SharedFrame:
    def __init__(self):
        self.lock = threading.Lock()
        self.image = None
        self.window = None
        self.boxes = {}
        self.frame_index = 0
        self.latest_detection = {"events": [], "badge_events": [], "change_events": [], "timestamp": utc_now()}
        self.latest_llm = None
        self.stop = threading.Event()

    def update_frame(self, image, window, boxes, frame_index):
        with self.lock:
            self.image = image
            self.window = window
            self.boxes = boxes
            self.frame_index = frame_index

    def snapshot(self):
        with self.lock:
            return self.image, self.window, dict(self.boxes), self.frame_index, dict(self.latest_detection), self.latest_llm

    def update_detection(self, detection):
        with self.lock:
            self.latest_detection = detection

    def update_llm(self, reading):
        with self.lock:
            self.latest_llm = reading


def parse_args():
    parser = argparse.ArgumentParser(description="Pure algorithm WeChat new-message monitor.")
    parser.add_argument("--regions", default=str(DEFAULT_REGIONS), help="wechat.json annotation path")
    parser.add_argument("--fps", type=float, default=30.0, help="preview FPS target")
    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 browser preview")
    parser.add_argument("--chat-diff-threshold", type=int, default=34, help="chat content pixel diff threshold")
    parser.add_argument("--dry-run", action="store_true", default=True, help="draw only; never click")
    parser.add_argument("--click-on-badge", action="store_true", help="click the first detected unread contact row")
    parser.add_argument("--llm-on-click", action="store_true", help="after clicking a contact, ask LLM to read current chat content")
    parser.add_argument("--llm-on-chat-change", action="store_true", help="ask LLM to read current chat when chat_content diff indicates a new message")
    parser.add_argument("--reply-on-chat-change", action="store_true", help="send reply_text for current-chat diff events")
    parser.add_argument("--send-reply", action="store_true", help="paste LLM reply_text into input_box and press Enter")
    parser.add_argument("--typing-chunk-size", type=int, default=2, help="characters pasted per simulated typing chunk")
    parser.add_argument("--typing-delay", type=float, default=0.08, help="seconds between simulated typing chunks")
    parser.add_argument("--agent-config", default=str(DEFAULT_AGENT_CONFIG), help="agent/config.toml path for LLM settings")
    parser.add_argument("--click-cooldown", type=float, default=30.0, help="seconds before clicking the same contact row again")
    parser.add_argument("--chat-change-cooldown", type=float, default=30.0, help="seconds before handling current chat diff again")
    parser.add_argument("--after-click-wait", type=float, default=1.0, help="seconds to wait before capturing chat after click")
    return parser.parse_args()


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


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


def emit_llm_reading(reading):
    payload = {
        "type": "llm_chat_reading_print",
        "timestamp": utc_now(),
        "contact_name": reading.get("contact_name", "unknown") if isinstance(reading, dict) else "unknown",
        "latest_user_message": reading.get("latest_user_message", "") if isinstance(reading, dict) else "",
        "visible_messages": reading.get("visible_messages", []) if isinstance(reading, dict) else [],
        "summary": reading.get("summary", "") if isinstance(reading, dict) else "",
        "reply_text": reading.get("reply_text", "") if isinstance(reading, dict) else "",
        "send_status": reading.get("send_status", "not_sent") if isinstance(reading, dict) else "not_sent",
        "send_error": reading.get("send_error", "") if isinstance(reading, dict) else "",
    }
    emit(payload)


def json_safe(value):
    if isinstance(value, dict):
        return {key: json_safe(item) for key, item in value.items()}
    if isinstance(value, list):
        return [json_safe(item) for item in value]
    if isinstance(value, tuple):
        return [json_safe(item) for item in value]
    if isinstance(value, np.integer):
        return int(value)
    if isinstance(value, np.floating):
        return float(value)
    return value


def load_regions(path):
    data = json.loads(Path(path).read_text(encoding="utf-8"))
    regions = {}
    for region in data.get("regions", []):
        regions.setdefault(region.get("type"), region)
        if region.get("type") == "custom" and "未读消息" in region.get("description", ""):
            regions["message_button"] = region
    return data, regions


def load_agent_config(path):
    with Path(path).open("rb") as file:
        return tomllib.load(file)


def image_data_url(image):
    buffer = BytesIO()
    image.save(buffer, format="JPEG", quality=88)
    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
    return "data:image/jpeg;base64," + encoded


def request_chat_reading(config, chat_image):
    volcengine = config.get("volcengine", {})
    base_url = (volcengine.get("base_url") or "").rstrip("/")
    api_key = volcengine.get("api_key") or ""
    model = volcengine.get("model") or ""
    if not base_url or not api_key or not model:
        raise RuntimeError("missing volcengine base_url/api_key/model in agent config")

    prompt = """你是微信聊天截图读取和回复草稿助手。
请读取图片中当前聊天区域的可见消息，不要编造看不见的内容。
请基于可见上下文生成一条自然、简短、适合直接发送的中文回复草稿。
这里只生成草稿，不要假设已经发送，也不要执行任何操作。
返回合法 JSON，格式如下：
{
  "contact_name": "unknown",
  "latest_user_message": "",
  "visible_messages": [
    {"sender":"unknown", "role":"user|me|system|unknown", "content":""}
  ],
  "summary": "",
  "reply_text": ""
}
"""
    payload = {
        "model": model,
        "temperature": 0.2,
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": image_data_url(chat_image)}},
                ],
            }
        ],
    }
    data = json.dumps(payload).encode("utf-8")
    req = urlrequest.Request(
        base_url + "/chat/completions",
        data=data,
        headers={"Authorization": "Bearer " + api_key, "Content-Type": "application/json"},
        method="POST",
    )
    with urlrequest.urlopen(req, timeout=90) as response:
        body = json.loads(response.read().decode("utf-8"))
    choices = body.get("choices") or []
    if not choices:
        raise RuntimeError("LLM returned no choices")
    content = choices[0].get("message", {}).get("content", "")
    return extract_json_object(content)


def extract_json_object(text):
    stripped = text.strip()
    if stripped.startswith("```json"):
        stripped = stripped[len("```json"):].strip()
    if stripped.startswith("```"):
        stripped = stripped[3:].strip()
    if stripped.endswith("```"):
        stripped = stripped[:-3].strip()
    start = stripped.find("{")
    end = stripped.rfind("}")
    if start >= 0 and end > start:
        stripped = stripped[start:end + 1]
    try:
        return json.loads(stripped)
    except json.JSONDecodeError:
        return {"raw_text": text[:1200]}


def scaled_box(region, annotation, image_size, window):
    # Annotation bbox_image is stored in physical pixels at annotation-time
    # scaleFactor. Convert back to logical window coords, then map to the
    # current capture's physical pixels. This stays accurate when the WeChat
    # window is moved or resized horizontally.
    annotation_scale = region.get("scaleFactor") or annotation.get("scaleFactor") or 1
    current_scale_x = image_size[0] / max(window["width"], 1)
    current_scale_y = image_size[1] / max(window["height"], 1)
    x1, y1, x2, y2 = region["bbox_image"]
    return [
        int((x1 / annotation_scale) * current_scale_x),
        int((y1 / annotation_scale) * current_scale_y),
        int((x2 / annotation_scale) * current_scale_x),
        int((y2 / annotation_scale) * current_scale_y),
    ]


def clamp_box(box, image_size):
    width, height = image_size
    x1, y1, x2, y2 = box
    x1, x2 = sorted((x1, x2))
    y1, y2 = sorted((y1, y2))
    return [max(0, x1), max(0, y1), min(width, x2), min(height, y2)]


def valid_box(box):
    return box and len(box) == 4 and box[2] > box[0] and box[3] > box[1]


def components(mask, min_area=12):
    h, w = mask.shape
    visited = np.zeros_like(mask, dtype=bool)
    found = []
    ys, xs = np.where(mask)
    for start_x, start_y in zip(xs, ys):
        if visited[start_y, start_x]:
            continue
        stack = [(start_x, start_y)]
        visited[start_y, start_x] = True
        area = 0
        min_x = max_x = start_x
        min_y = max_y = start_y
        while stack:
            x, y = stack.pop()
            area += 1
            min_x, max_x = min(min_x, x), max(max_x, x)
            min_y, max_y = min(min_y, y), max(max_y, y)
            for nx, ny in ((x - 1, y), (x + 1, y), (x, y - 1), (x, y + 1)):
                if 0 <= nx < w and 0 <= ny < h and mask[ny, nx] and not visited[ny, nx]:
                    visited[ny, nx] = True
                    stack.append((nx, ny))
        if area >= min_area:
            found.append({"bbox": [min_x, min_y, max_x + 1, max_y + 1], "area": area})
    return found


def detect_red_badges(image, contact_box):
    x1, y1, x2, y2 = contact_box
    crop = np.asarray(image.crop((x1, y1, x2, y2)).convert("RGB"))
    r, g, b = crop[:, :, 0], crop[:, :, 1], crop[:, :, 2]
    mask = (r > 175) & (g < 105) & (b < 105) & ((r.astype(int) - g.astype(int)) > 65) & ((r.astype(int) - b.astype(int)) > 65)
    digit_mask = (r > 215) & (g > 215) & (b > 215) & (np.maximum.reduce([r, g, b]) - np.minimum.reduce([r, g, b]) < 42)
    results = []
    for comp in components(mask, min_area=18):
        bx1, by1, bx2, by2 = comp["bbox"]
        bw, bh = bx2 - bx1, by2 - by1
        if not (6 <= bw <= 52 and 6 <= bh <= 52):
            continue
        if comp["area"] / max(bw * bh, 1) < 0.22:
            continue
        pad_x = max(2, bw // 5)
        pad_y = max(2, bh // 5)
        inner_x1 = min(bx2, bx1 + pad_x)
        inner_y1 = min(by2, by1 + pad_y)
        inner_x2 = max(inner_x1, bx2 - pad_x)
        inner_y2 = max(inner_y1, by2 - pad_y)
        digit_pixels = int(digit_mask[inner_y1:inner_y2, inner_x1:inner_x2].sum())
        # Require white digit strokes inside the red badge. This intentionally
        # ignores plain red dots and unrelated red UI fragments.
        if digit_pixels < 3:
            continue
        gx1, gy1, gx2, gy2 = x1 + bx1, y1 + by1, x1 + bx2, y1 + by2
        center_y = (gy1 + gy2) // 2
        # WeChat conversation rows are roughly 64-76 logical px tall. The
        # captured image is usually Retina scale 2, so 132 image px is a good
        # first approximation and less noisy than deriving height from badge.
        row_h = 132
        row = [x1, max(y1, center_y - row_h // 2), x2, min(y2, center_y + row_h // 2)]
        click = [(row[0] + row[2]) // 2, (row[1] + row[2] * 0 + row[3]) // 2]
        results.append({"type": "new_message_badge", "badge_box": [gx1, gy1, gx2, gy2], "contact_row": row, "click_image": click, "score": round(float(comp["area"] / max(bw * bh, 1)), 3), "digit_pixels": digit_pixels})
    return results


def detect_chat_change(image, previous_gray, chat_box, threshold):
    x1, y1, x2, y2 = chat_box
    gray = np.asarray(image.crop((x1, y1, x2, y2)).convert("L"), dtype=np.int16)
    if previous_gray is None or previous_gray.shape != gray.shape:
        return gray, []
    diff = np.abs(gray - previous_gray)
    mask = diff > threshold
    crop_h, crop_w = gray.shape
    changed_ratio = float(mask.sum()) / max(crop_w * crop_h, 1)
    if changed_ratio > 0.28:
        # Window moves, scrolls, or large reflows are not new-message events.
        return gray, []
    results = []
    for comp in components(mask, min_area=120):
        bx1, by1, bx2, by2 = comp["bbox"]
        bw, bh = bx2 - bx1, by2 - by1
        if bw < 24 or bh < 10:
            continue
        if by1 < crop_h * 0.38:
            continue
        if comp["area"] > crop_w * crop_h * 0.18:
            continue
        # Focus on localized lower-area changes where new messages appear.
        gx1, gy1, gx2, gy2 = x1 + bx1, y1 + by1, x1 + bx2, y1 + by2
        results.append({"type": "current_chat_new_message", "change_box": [gx1, gy1, gx2, gy2], "area": comp["area"], "changed_ratio": round(changed_ratio, 4)})
    return gray, results[:5]


def image_to_screen_box(box, window, image_size):
    scale_x = image_size[0] / max(window["width"], 1)
    scale_y = image_size[1] / max(window["height"], 1)
    x1, y1, x2, y2 = box
    return [round(window["x"] + x1 / scale_x, 2), round(window["y"] + y1 / scale_y, 2), round(window["x"] + x2 / scale_x, 2), round(window["y"] + y2 / scale_y, 2)]


def point_to_screen(point, window, image_size):
    scale_x = image_size[0] / max(window["width"], 1)
    scale_y = image_size[1] / max(window["height"], 1)
    return [round(window["x"] + point[0] / scale_x, 2), round(window["y"] + point[1] / scale_y, 2)]


def box_center_screen(box, window, image_size):
    cx = (box[0] + box[2]) / 2
    cy = (box[1] + box[3]) / 2
    return point_to_screen([cx, cy], window, image_size)


def text_chunks(text, chunk_size):
    runes = list(text)
    size = max(1, chunk_size)
    for index in range(0, len(runes), size):
        yield "".join(runes[index:index + size])


def send_reply_text(reply_text, input_box, window, image_size, chunk_size=2, typing_delay=0.08):
    text = (reply_text or "").strip()
    if not text:
        return {"send_status": "skipped", "send_error": "empty reply_text"}
    if not valid_box(input_box):
        return {"send_status": "failed", "send_error": "invalid input_box"}
    point = box_center_screen(input_box, window, image_size)
    pyautogui.click(point[0], point[1])
    time.sleep(0.15)
    chunk_count = 0
    for chunk in text_chunks(text, chunk_size):
        pyperclip.copy(chunk)
        pyautogui.hotkey("command", "v")
        chunk_count += 1
        time.sleep(max(0.01, typing_delay))
    pyautogui.press("enter")
    return {"send_status": "sent", "send_error": "", "input_click_screen": point, "typing_mode": "chunked_paste", "typing_chunk_size": max(1, chunk_size), "typing_chunks": chunk_count}


def decorate_events(events, window, image_size):
    decorated = []
    for event in events:
        item = dict(event)
        for key in ("badge_box", "contact_row", "change_box"):
            if key in item:
                item[key + "_screen"] = image_to_screen_box(item[key], window, image_size)
        if "click_image" in item:
            item["click_screen"] = point_to_screen(item["click_image"], window, image_size)
        item["dry_run"] = True
        decorated.append(item)
    return decorated


def draw_events(image, events, boxes):
    draw = ImageDraw.Draw(image)
    for name, box in boxes.items():
        if not valid_box(box):
            continue
        color = {"contact_list": "#666666", "chat_content": "#5555ff", "input_box": "#22c55e", "message_button": "#888888"}.get(name, "#666666")
        draw.rectangle(box, outline=color, width=2)
        draw.text((box[0] + 4, box[1] + 4), name, fill=color)
    for event in events:
        if event["type"] == "new_message_badge":
            if valid_box(event.get("badge_box")):
                draw.rectangle(event["badge_box"], outline="red", width=4)
            if valid_box(event.get("contact_row")):
                draw.rectangle(event["contact_row"], outline="yellow", width=4)
                draw.text((event["contact_row"][0] + 8, event["contact_row"][1] + 8), "new message", fill="yellow")
        elif event["type"] in {"chat_content_changed", "current_chat_new_message"}:
            if valid_box(event.get("change_box")):
                draw.rectangle(event["change_box"], outline="#00aaff", width=3)
                draw.text((event["change_box"][0] + 8, event["change_box"][1] + 8), "chat diff", fill="#00aaff")
    return image


def draw_overlay(image, event):
    draw = ImageDraw.Draw(image)
    text = f"frame={event.get('frame_index')} badges={event.get('badge_count')} changes={event.get('change_count')} click={event.get('click_enabled')} llm={event.get('llm_enabled')}"
    draw.rectangle((12, 12, 980, 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 Monitor</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)460px;gap:16px;padding:16px}
    .stage{background:#18181b;border:1px solid #27272a;border-radius:14px;overflow:hidden;align-self:start}
    img{display:block;width:100%;height:auto}
    aside{background:#18181b;border:1px solid #27272a;border-radius:14px;padding:16px;max-height:calc(100vh - 32px);overflow:auto}
    h1{margin:0 0 12px;font-size:18px}.item{padding:12px 0;border-top:1px solid #27272a}.label{color:#a1a1aa;font-size:12px;margin-bottom:6px}
    .value{font-family:ui-monospace,SFMono-Regular,Menlo,monospace;white-space:pre-wrap;word-break:break-all;font-size:12px}
    .msg{padding:8px 10px;border-radius:10px;margin:8px 0;background:#27272a}.msg.user{background:#1e3a8a}.msg.me{background:#14532d}.role{font-size:11px;color:#d4d4d8;margin-bottom:4px}.content{font-size:14px;line-height:1.45}
    .reply{padding:12px;border-radius:12px;background:#052e16;color:#bbf7d0;font-size:15px;line-height:1.5}.empty{color:#71717a}.ok{color:#86efac}.warn{color:#facc15}
    @media(max-width:1000px){.wrap{grid-template-columns:1fr}aside{max-height:none}}
  </style>
</head>
<body>
  <div class="wrap">
    <main class="stage"><img src="/stream.mjpg" alt="live stream"></main>
    <aside>
      <h1>WeChat Monitor</h1>
      <div class="item"><div class="label">Sequence / Status</div><div id="status" class="value">-</div></div>
      <div class="item"><div class="label">监测事件</div><div id="eventSummary" class="value">-</div></div>
      <div class="item"><div class="label">最新用户消息</div><div id="latest" class="value empty">等待 LLM 读取...</div></div>
      <div class="item"><div class="label">识别到的聊天记录</div><div id="messages"></div></div>
      <div class="item"><div class="label">建议回复（未发送）</div><div id="reply" class="reply empty">等待生成...</div></div>
      <div class="item"><div class="label">摘要</div><div id="summary" class="value empty">-</div></div>
    </aside>
  </div>
  <script>
    function escapeHtml(text){return String(text ?? '').replace(/[&<>'"]/g,c=>({'&':'&amp;','<':'&lt;','>':'&gt;',"'":'&#39;','"':'&quot;'}[c]))}
    function messageHtml(message){
      const role = message.role || 'unknown';
      const cls = role === 'user' ? 'user' : (role === 'me' ? 'me' : '');
      const sender = message.sender || role;
      return `<div class="msg ${cls}"><div class="role">${escapeHtml(sender)} · ${escapeHtml(role)}</div><div class="content">${escapeHtml(message.content || '')}</div></div>`;
    }
    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 llm = event.latest_llm || event.llm_chat_reading || null;
        document.getElementById('status').textContent = `seq=${data.sequence ?? '-'} frame=${event.frame_index ?? '-'} badges=${event.badge_count ?? 0} changes=${event.change_count ?? 0}`;
        document.getElementById('eventSummary').textContent = JSON.stringify({type:event.type, click_enabled:event.click_enabled, llm_enabled:event.llm_enabled, events:event.events || []}, null, 2);
        if(llm){
          document.getElementById('latest').textContent = llm.latest_user_message || '';
          document.getElementById('latest').className = 'value';
          const messages = Array.isArray(llm.visible_messages) ? llm.visible_messages : [];
          document.getElementById('messages').innerHTML = messages.length ? messages.map(messageHtml).join('') : '<div class="empty">未识别到可见消息</div>';
          const reply = llm.reply_text || '';
          document.getElementById('reply').textContent = reply || '未生成建议回复';
          document.getElementById('reply').className = 'reply' + (reply ? '' : ' empty');
          document.getElementById('summary').textContent = llm.summary || '-';
          document.getElementById('summary').className = 'value';
        }
      }catch(error){
        document.getElementById('status').textContent = '连接中断或服务未启动';
      }
    }
    refresh();
    setInterval(refresh, 500);
  </script>
</body>
</html>"""


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

        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("Content-Length", str(len(data)))
                self.end_headers()
                self.wfile.write(data)
                return
            if self.path.startswith("/latest.json"):
                _, event, sequence = state.snapshot()
                data = json.dumps(json_safe({"sequence": sequence, "updated_at": utc_now(), "event": event or {"type": "waiting"}}), 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.sequence
                    if not jpeg:
                        continue
                    try:
                        self.wfile.write(b"--frame\r\nContent-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))
    threading.Thread(target=server.serve_forever, daemon=True).start()
    return server


def boxes_for_frame(annotation, regions, image, window):
    boxes = {name: clamp_box(scaled_box(regions[name], annotation, image.size, window), image.size) for name in ("contact_list", "chat_content")}
    if "input_box" in regions:
        boxes["input_box"] = clamp_box(scaled_box(regions["input_box"], annotation, image.size, window), image.size)
    if "message_button" in regions:
        boxes["message_button"] = clamp_box(scaled_box(regions["message_button"], annotation, image.size, window), image.size)
    return boxes


def preview_loop(args, annotation, regions, stream_state, shared):
    interval = 1 / max(args.fps, 1)
    frame_index = 0
    while not shared.stop.is_set() and (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()}
            stream_state.update(image_to_jpeg_bytes(Image.new("RGB", (960, 540), "black")), event)
            time.sleep(interval)
            continue

        image = capture_window_image(window["id"])
        if image is None:
            time.sleep(interval)
            continue

        frame_index += 1
        boxes = boxes_for_frame(annotation, regions, image, window)
        shared.update_frame(image, window, boxes, frame_index)
        _, _, _, _, detection, latest_llm = shared.snapshot()
        badge_events = detection.get("badge_events", [])
        change_events = detection.get("change_events", [])
        event = {
            "type": "algorithm_monitor",
            "timestamp": utc_now(),
            "frame_index": frame_index,
            "badge_count": len(badge_events),
            "change_count": len(change_events),
            "events": detection.get("events", []),
            "window": window,
            "dry_run": not args.click_on_badge,
            "click_enabled": args.click_on_badge,
            "llm_enabled": args.llm_on_click,
            "preview_only": True,
        }
        if latest_llm:
            event["latest_llm"] = latest_llm
        annotated = draw_events(image.copy(), badge_events + change_events, boxes)
        annotated = draw_overlay(annotated, event)
        stream_state.update(image_to_jpeg_bytes(annotated), event)
        time.sleep(max(0, interval - (time.time() - start)))
    shared.stop.set()


def llm_worker(args, annotation, regions, agent_config, shared, job_event):
    last_click_at = {}
    last_chat_change_at = 0.0
    while not shared.stop.is_set():
        if not job_event.wait(timeout=0.5):
            continue
        job_event.clear()
        image, window, boxes, frame_index, detection, _ = shared.snapshot()
        if image is None or window is None:
            continue
        target = next((item for item in detection.get("events", []) if item.get("type") == "new_message_badge"), None)
        chat_change = next((item for item in detection.get("events", []) if item.get("type") == "current_chat_new_message"), None)
        trigger = "unread_badge" if target else ("current_chat_change" if chat_change else "")
        if not trigger:
            continue

        if trigger == "unread_badge":
            click_key = str(round(target["click_screen"][1] / 10) * 10)
            if time.time() - last_click_at.get(click_key, 0) < args.click_cooldown:
                continue
            if args.click_on_badge:
                pyautogui.click(target["click_screen"][0], target["click_screen"][1])
                last_click_at[click_key] = time.time()
            else:
                continue
            if not args.llm_on_click:
                continue
        elif trigger == "current_chat_change":
            if not args.llm_on_chat_change:
                continue
            if time.time() - last_chat_change_at < args.chat_change_cooldown:
                continue
            last_chat_change_at = time.time()

        time.sleep(args.after_click_wait)
        refreshed_window = find_wechat_window() or window
        refreshed_image = capture_window_image(refreshed_window["id"])
        if refreshed_image is None:
            continue
        refreshed_chat_box = clamp_box(scaled_box(regions["chat_content"], annotation, refreshed_image.size, refreshed_window), refreshed_image.size)
        if not valid_box(refreshed_chat_box):
            continue
        x1, y1, x2, y2 = refreshed_chat_box
        chat_crop = refreshed_image.crop((x1, y1, x2, y2))
        try:
            reading = request_chat_reading(agent_config, chat_crop)
            reading["send_status"] = "not_sent"
            reading["trigger"] = trigger
            should_send = args.send_reply and (trigger == "unread_badge" or args.reply_on_chat_change)
            if should_send:
                input_region = regions.get("input_box")
                if input_region:
                    input_box = clamp_box(scaled_box(input_region, annotation, refreshed_image.size, refreshed_window), refreshed_image.size)
                    send_result = send_reply_text(
                        reading.get("reply_text", ""),
                        input_box,
                        refreshed_window,
                        refreshed_image.size,
                        args.typing_chunk_size,
                        args.typing_delay,
                    )
                    reading.update(send_result)
                else:
                    reading["send_status"] = "failed"
                    reading["send_error"] = "missing input_box region"
            shared.update_llm(reading)
            emit_llm_reading(reading)
        except Exception as error:
            emit({"type": "llm_error", "timestamp": utc_now(), "error": str(error)})


def detection_loop(args, shared, job_event):
    previous_chat_gray = None
    interval = 1.0
    while not shared.stop.is_set():
        start = time.time()
        image, window, boxes, frame_index, _, _ = shared.snapshot()
        if image is not None and window is not None and "contact_list" in boxes and "chat_content" in boxes:
            badge_events = detect_red_badges(image, boxes["contact_list"])
            previous_chat_gray, change_events = detect_chat_change(image, previous_chat_gray, boxes["chat_content"], args.chat_diff_threshold)
            events = decorate_events(badge_events + change_events, window, image.size)
            detection = {
                "timestamp": utc_now(),
                "frame_index": frame_index,
                "events": events,
                "badge_events": badge_events,
                "change_events": change_events,
            }
            shared.update_detection(detection)
            event = {
                "type": "algorithm_detection",
                "timestamp": utc_now(),
                "frame_index": frame_index,
                "badge_count": len(badge_events),
                "change_count": len(change_events),
                "events": events,
                "window": window,
                "detect_interval_sec": interval,
            }
            if badge_events or change_events:
                emit(event)
            if badge_events and args.click_on_badge:
                job_event.set()
            elif change_events and args.llm_on_chat_change:
                job_event.set()
        time.sleep(max(0, interval - (time.time() - start)))


def run(args):
    annotation, regions = load_regions(args.regions)
    agent_config = load_agent_config(args.agent_config) if args.llm_on_click else None
    required = ["contact_list", "chat_content"]
    missing = [name for name in required if name not in regions]
    if missing:
        raise RuntimeError(f"missing required regions: {missing}")
    state = StreamState()
    server = start_server(args.host, args.port, state)
    url = f"http://{args.host}:{args.port}/"
    emit({"type": "server_started", "url": url, "optimized": True})
    if args.open_browser:
        webbrowser.open(url)

    shared = SharedFrame()
    job_event = threading.Event()
    threads = [
        threading.Thread(target=preview_loop, args=(args, annotation, regions, state, shared), daemon=True),
        threading.Thread(target=detection_loop, args=(args, shared, job_event), daemon=True),
    ]
    if args.click_on_badge and args.llm_on_click:
        threads.append(threading.Thread(target=llm_worker, args=(args, annotation, regions, agent_config, shared, job_event), daemon=True))
    for thread in threads:
        thread.start()
    try:
        while not shared.stop.is_set():
            time.sleep(0.5)
    finally:
        shared.stop.set()
        server.shutdown()


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()
