"""Serve the System-2 backbone (Qwen3.5-9B) as an OpenAI-compatible endpoint.

Preferred path - vLLM (recommended by the Qwen3.5 recipes, single GPU):
    vllm serve Qwen/Qwen3.5-9B --port 8000

This script is the fallback when vLLM is unavailable: it wraps transformers
in a minimal /v1/chat/completions server. Both expose the same contract the
System2Controller expects (mode="openai").
"""

from __future__ import annotations

import argparse
import base64
import io
import re


def extract_messages(payload: dict) -> tuple[list[dict], dict]:
    messages = payload.get("messages", [])
    gen_kwargs = {
        "temperature": payload.get("temperature", 0.2),
        "max_new_tokens": payload.get("max_tokens", 2048),
    }
    return messages, gen_kwargs


def decode_images(messages: list[dict]) -> list:
    """Pull base64 PNGs out of OpenAI-style image_url content blocks."""
    images = []
    for msg in messages:
        content = msg.get("content")
        if isinstance(content, list):
            for part in content:
                if part.get("type") == "image_url":
                    url = part["image_url"]["url"]
                    b64 = re.sub("^data:image/\\w+;base64,", "", url)
                    from PIL import Image

                    images.append(Image.open(io.BytesIO(base64.b64decode(b64))))
    return images


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="Qwen/Qwen3.5-9B")
    ap.add_argument("--port", type=int, default=8000)
    args = ap.parse_args()

    import torch
    from flask import Flask, jsonify, request
    from transformers import AutoProcessor, AutoModelForImageTextToText

    processor = AutoProcessor.from_pretrained(args.model)
    model = AutoModelForImageTextToText.from_pretrained(
        args.model, torch_dtype=torch.bfloat16, device_map="auto"
    )

    app = Flask(__name__)

    @app.post("/v1/chat/completions")
    def completions():
        payload = request.get_json(force=True)
        messages, gen = extract_messages(payload)
        text = processor.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=False
        )
        inputs = processor(text=text, images=decode_images(messages) or None,
                           return_tensors="pt").to(model.device)
        out = model.generate(**inputs, **gen)
        completion = processor.batch_decode(out[:, inputs["input_ids"].shape[1]:],
                                            skip_special_tokens=True)[0]
        return jsonify({
            "id": "qwenjev-s2",
            "choices": [{"message": {"role": "assistant", "content": completion},
                         "finish_reason": "stop"}],
        })

    app.run(host="127.0.0.1", port=args.port)


if __name__ == "__main__":
    main()