import argparse
import base64
import io
import json
import sys
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import torch
from PIL import Image
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration


ALLOWED_TAGS = {
    "镜头单人",
    "双人分屏",
    "多人分屏",
    "画面清晰",
    "明显展示动作",
    "镜头互动",
    "画面静止",
    "可见道具展示",
}

PROMPT = """只分析直播画面中明确可见的布局和表现，不判断身份、性别、年龄、情绪、关系、消费或直播成绩。
只输出一个 JSON 对象，不要 Markdown，不要解释：
{"layout":"solo|split|four_panel|unknown","pk_ui_visible":true|false|null,"visual_tags":["从允许列表选择，最多3个"],"description":"不超过24字的客观画面描述"}
visual_tags 只能从以下词中选择：镜头单人、双人分屏、多人分屏、画面清晰、明显展示动作、镜头互动、画面静止、可见道具展示。
看不清就用 unknown 或 null；pk_ui_visible 只有清楚看见 PK 界面时才 true。"""


PK_READOUT_PROMPT = """You receive one composite of five enlarged crops from the same livestream screenshot. The timer crop is top-centre. The four badge crops are arranged in their original positions: top-left, top-right, bottom-left, bottom-right.
Read only digits that are visibly present: the PK countdown, rank, and score/heat value. Never infer a number from PK rules, people, names, or prior frames. If a digit is unclear, blocked, or absent, return null.
Convert a visible mm:ss timer to remaining_seconds; for example 02:23 becomes 143. Use positions top_left, top_right, bottom_left, bottom_right.
Return one JSON object only, without Markdown or explanation:
{"pk_visible":true|false|null,"remaining_seconds":integer|null,"tiles":[{"position":"top_left","rank":integer|null,"value":integer|null},{"position":"top_right","rank":integer|null,"value":integer|null},{"position":"bottom_left","rank":integer|null,"value":integer|null},{"position":"bottom_right","rank":integer|null,"value":integer|null}]}
Keep a tile but use null for rank/value whenever it is not reliably readable."""


def extract_json(text):
    clean = (text or "").strip()
    if clean.startswith("```"):
        clean = clean.split("\n", 1)[-1]
        if clean.rstrip().endswith("```"):
            clean = clean.rstrip()[:-3]
    start = clean.find("{")
    end = clean.rfind("}")
    if start < 0 or end <= start:
        return {}
    try:
        return json.loads(clean[start:end + 1])
    except Exception:
        return {}


class VisionModel:
    def __init__(self, model_dir):
        print("VISION_MODEL_LOADING", flush=True)
        self.model = Qwen3VLForConditionalGeneration.from_pretrained(
            model_dir,
            torch_dtype=torch.float16,
            device_map="auto",
            local_files_only=True,
        )
        self.processor = AutoProcessor.from_pretrained(model_dir, local_files_only=True)
        self.lock = threading.Lock()
        self.device = str(self.model.device)
        print(f"VISION_MODEL_READY|device={self.device}", flush=True)

    def analyze(self, image_bytes):
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": PROMPT},
            ],
        }]
        with self.lock, torch.inference_mode():
            inputs = self.processor.apply_chat_template(
                messages,
                tokenize=True,
                add_generation_prompt=True,
                return_dict=True,
                return_tensors="pt",
            ).to(self.model.device)
            output = self.model.generate(**inputs, max_new_tokens=112, do_sample=False)
            generated = output[:, inputs.input_ids.shape[1]:]
            raw = self.processor.batch_decode(generated, skip_special_tokens=True)[0]

        parsed = extract_json(raw)
        layout = parsed.get("layout")
        if layout not in {"solo", "split", "four_panel"}:
            layout = "unknown"
        pk_ui = parsed.get("pk_ui_visible")
        if not isinstance(pk_ui, bool):
            pk_ui = None
        tags = parsed.get("visual_tags")
        if not isinstance(tags, list):
            tags = []
        tags = [str(tag).strip() for tag in tags if str(tag).strip() in ALLOWED_TAGS][:3]
        description = str(parsed.get("description") or "").replace("\n", " ").strip()[:48]
        return {
            "ok": True,
            "layout": layout,
            "pk_ui_visible": pk_ui,
            "visual_tags": tags,
            "description": description,
            "confidence": 1.0 if bool(parsed) else 0.0,
            "raw_valid": bool(parsed),
        }

    def analyze_pk_readout(self, image_bytes):
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        content = [
            {"type": "image", "image": self._pk_readout_composite(image)},
            {"type": "text", "text": PK_READOUT_PROMPT},
        ]
        messages = [{"role": "user", "content": content}]
        with self.lock, torch.inference_mode():
            inputs = self.processor.apply_chat_template(
                messages,
                tokenize=True,
                add_generation_prompt=True,
                return_dict=True,
                return_tensors="pt",
            ).to(self.model.device)
            output = self.model.generate(**inputs, max_new_tokens=220, do_sample=False)
            generated = output[:, inputs.input_ids.shape[1]:]
            raw = self.processor.batch_decode(generated, skip_special_tokens=True)[0]

        parsed = extract_json(raw)
        tiles = parsed.get("tiles")
        if not isinstance(tiles, list):
            tiles = []
        return {
            "ok": True,
            "pk_visible": parsed.get("pk_visible") if isinstance(parsed.get("pk_visible"), bool) else None,
            "remaining_seconds": self._safe_int(parsed.get("remaining_seconds"), 0, 1200),
            "tiles": [self._normalize_tile(tile) for tile in tiles if isinstance(tile, dict)][:4],
            "confidence": 1.0 if bool(parsed) else 0.0,
            "raw_valid": bool(parsed),
        }

    @staticmethod
    def _safe_int(value, lower, upper):
        if isinstance(value, bool):
            return None
        try:
            number = int(value)
        except (TypeError, ValueError):
            return None
        return number if lower <= number <= upper else None

    def _normalize_tile(self, tile):
        position = str(tile.get("position") or "").strip()
        if position not in {"top_left", "top_right", "bottom_left", "bottom_right"}:
            return {"position": "unknown", "rank": None, "value": None}
        return {
            "position": position,
            "rank": self._safe_int(tile.get("rank"), 1, 4),
            "value": self._safe_int(tile.get("value"), 0, 9999999),
        }

    @staticmethod
    def _pk_readout_composite(image):
        width, height = image.size
        specs = [
            ((0.35, 0.00, 0.65, 0.13), (300, 0, 360, 150)),
            ((0.00, 0.00, 0.25, 0.16), (0, 160, 300, 190)),
            ((0.75, 0.00, 1.00, 0.16), (660, 160, 300, 190)),
            ((0.00, 0.46, 0.25, 0.66), (0, 360, 300, 190)),
            ((0.75, 0.46, 1.00, 0.66), (660, 360, 300, 190)),
        ]
        canvas = Image.new("RGB", (960, 560), (18, 18, 18))
        for (left, top, right, bottom), (x, y, target_width, target_height) in specs:
            box = (
                max(0, int(width * left)),
                max(0, int(height * top)),
                min(width, int(width * right)),
                min(height, int(height * bottom)),
            )
            crop = image.crop(box).resize((target_width, target_height), Image.Resampling.LANCZOS)
            canvas.paste(crop, (x, y))
        return canvas


def make_handler(model):
    class Handler(BaseHTTPRequestHandler):
        def log_message(self, *_):
            return

        def send_json(self, status, payload):
            body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
            self.send_response(status)
            self.send_header("Content-Type", "application/json; charset=utf-8")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self):
            if self.path == "/health":
                self.send_json(200, {"status": "ready", "device": model.device})
                return
            self.send_json(404, {"error": "not_found"})

        def do_POST(self):
            if self.path not in {"/analyze", "/analyze-pk"}:
                self.send_json(404, {"error": "not_found"})
                return
            try:
                length = int(self.headers.get("Content-Length", "0"))
                if length <= 0 or length > 12 * 1024 * 1024:
                    raise ValueError("invalid_image_size")
                payload = json.loads(self.rfile.read(length).decode("utf-8"))
                image_bytes = base64.b64decode(payload.get("image_base64", ""), validate=True)
                if len(image_bytes) == 0:
                    raise ValueError("empty_image")
                if self.path == "/analyze-pk":
                    self.send_json(200, model.analyze_pk_readout(image_bytes))
                else:
                    self.send_json(200, model.analyze(image_bytes))
            except Exception as error:
                self.send_json(400, {"error": str(error)[:160]})

    return Handler


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-dir", required=True)
    parser.add_argument("--port", required=True, type=int)
    args = parser.parse_args()
    model = VisionModel(args.model_dir)
    server = ThreadingHTTPServer(("127.0.0.1", args.port), make_handler(model))
    print(f"VISION_SERVER_READY|port={args.port}", flush=True)
    server.serve_forever()


if __name__ == "__main__":
    main()
