# https://huggingface.co/Qwen/Qwen3-4B
# https://github.com/QwenLM/Qwen3

import json
import logging
import os
import re
import threading
import time

import torch
from flask import Flask, abort, jsonify, request
from transformers import AutoModelForCausalLM, AutoTokenizer

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Device
# ---------------------------------------------------------------------------

RUN_ON_CPU = os.environ.get("RUN_ON_CPU", "0") == "1"
if RUN_ON_CPU:
    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

device = "cpu" if RUN_ON_CPU else ("cuda" if torch.cuda.is_available() else "cpu")
logger.info("Device: %s", device)

# Serialize model.generate() — not thread-safe
lock = threading.Lock()

# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------

MODEL_NAME = os.environ.get("LLM_MODEL", "Qwen/Qwen3-4B")

logger.info("Loading model: %s", MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype="auto",
    device_map="auto",
)
model.eval()
logger.info("Model loaded: %s  (max_position_embeddings=%d)",
            MODEL_NAME, model.config.max_position_embeddings)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

cache_lock = threading.Lock()
cache_cnt = 0

def empty_cache():
    global cache_cnt
    with cache_lock:
        cache_cnt += 1
        flush = (cache_cnt % 10 == 0)
    if flush and torch.cuda.is_available():
        torch.cuda.empty_cache()

# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

def strip_thinking(text: str) -> str:
    """Remove <think>…</think> block emitted by Qwen3 in thinking mode."""
    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()


def run(data: dict) -> str:
    task = (data.get("task") or "").strip()
    if not task:
        abort(400, description="Missing or empty 'task' field")

    max_new_tokens  = int(data.get("max_new_tokens") or data.get("max_tokens") or 500)
    temperature     = float(data.get("temperature", 0.15))
    top_p           = float(data.get("top_p", 0.9))
    top_k           = int(data.get("top_k", 20))
    rep_penalty     = float(data.get("repetition_penalty", 1.1))
    enable_thinking = bool(data.get("enable_thinking", False))
    system_prompt   = data.get("prompt") or data.get("system_prompt")

    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": task})

    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=enable_thinking,
    )

    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    input_len = model_inputs.input_ids.shape[1]

    logger.info("generate | input_tokens=%d max_new=%d temp=%.2f top_p=%.2f thinking=%s",
                input_len, max_new_tokens, temperature, top_p, enable_thinking)

    gen_kwargs = dict(
        **model_inputs,
        max_new_tokens=max_new_tokens,
        repetition_penalty=rep_penalty,
    )
    if temperature > 0:
        gen_kwargs.update(do_sample=True, temperature=temperature, top_p=top_p, top_k=top_k)
    else:
        gen_kwargs["do_sample"] = False

    with lock, torch.no_grad():
        generated_ids = model.generate(**gen_kwargs)

    output_ids = generated_ids[0][input_len:].tolist()
    result = tokenizer.decode(output_ids, skip_special_tokens=True).strip()

    if not enable_thinking:
        result = strip_thinking(result)

    logger.info("result (%d chars): %r", len(result), result[:300])
    return result

# ---------------------------------------------------------------------------
# Flask app
# ---------------------------------------------------------------------------

# http://localhost:8087/run?task=Who+are+you%3F

app = Flask(__name__)
app.config['DEBUG'] = False


@app.route("/run", methods=['GET'])
def run_get():
    task = (request.args.get('task') or '').strip()
    if not task:
        abort(400, description="Missing 'task' query parameter")
    try:
        empty_cache()
        start = time.perf_counter()
        result = run({"task": task, "max_new_tokens": 500})

        took = int((time.perf_counter() - start) * 1000)
        return jsonify({'result': result, 'took': took})
    except RuntimeError as e:
        logger.error("Generation error: %s", e)
        abort(500, description=str(e))


@app.route("/run", methods=['POST'])
def run_post():
    try:
        data = json.loads(request.data)
    except (json.JSONDecodeError, ValueError):
        abort(400, description="Invalid JSON body")
    try:
        empty_cache()
        start = time.perf_counter()
        result = run(data)
        took = int((time.perf_counter() - start) * 1000)
        return jsonify({'result': result, 'took': took})
    except RuntimeError as e:
        logger.error("Generation error: %s", e)
        abort(500, description=str(e))


@app.route("/health", methods=['GET'])
def health():
    return jsonify({"status": "ok", "model": MODEL_NAME, "device": device})


if __name__ == '__main__':
    app.run(host="0.0.0.0", debug=False)
