#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import logging
import os
import sys
import torch
from flask import Flask, jsonify, request
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)

app = Flask(__name__)

MODEL = None
TOKENIZER = None
DEVICE = None
STOP_ID = None

def load_model(model_dir: str, base_model_id: str, quantize: bool = False):
    global MODEL, TOKENIZER, DEVICE, STOP_ID

    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    log.info(f"Using device: {DEVICE}")

    if quantize:
        qcfg = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
        )
        MODEL = AutoModelForCausalLM.from_pretrained(
            model_dir,
            quantization_config=qcfg,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
        )
    else:
        MODEL = AutoModelForCausalLM.from_pretrained(
            model_dir,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
        )

    MODEL.eval().to(DEVICE)

    TOKENIZER = AutoTokenizer.from_pretrained(
        base_model_id,
        trust_remote_code=True,
        use_fast=False,
    )

    TOKENIZER.pad_token = TOKENIZER.eos_token

    if "<STOP>" not in TOKENIZER.get_vocab():
        log.info("Adding <STOP> token to tokenizer")
        TOKENIZER.add_special_tokens(
            {"additional_special_tokens": ["<STOP>"]}
        )
        MODEL.resize_token_embeddings(len(TOKENIZER))

    STOP_ID = TOKENIZER.convert_tokens_to_ids("<STOP>")
    log.info(f"<STOP> token id: {STOP_ID}")

    log.info("Model & tokenizer loaded successfully")

def build_prompt(text: str, source: str, target: str) -> str:
    if source == "en" and target == "es":
        return f"Translate the following English text to Spanish:\n{text}\n\nTranslation:"
    elif source == "es" and target == "en":
        return f"Translate the following Spanish text to English:\n{text}\n\nTranslation:"
    else:
        raise ValueError("Unsupported translation direction")

@torch.inference_mode()
def translate(text: str, source: str, target: str) -> str:
    prompt = build_prompt(text, source, target)

    inputs = TOKENIZER(prompt, return_tensors="pt").to(DEVICE)
    prompt_len = inputs["input_ids"].shape[1]

    src_tokens = len(TOKENIZER.tokenize(text))
    max_new = int(src_tokens * 1.3) + 6

    output = MODEL.generate(
        **inputs,
        max_new_tokens=max_new,
        do_sample=False,
        temperature=0.0,
        eos_token_id=STOP_ID,
        pad_token_id=TOKENIZER.eos_token_id,
        repetition_penalty=1.05,
    )

    decoded = TOKENIZER.decode(
        output[0][prompt_len:], skip_special_tokens=False
    )

    return decoded.split("<STOP>")[0].strip()

@app.route("/translate", methods=["POST"])
def translate_endpoint():
    data = request.get_json(silent=True)
    if not data:
        return jsonify({"error": "Invalid JSON"}), 400

    text = data.get("text")
    source = data.get("source")
    target = data.get("target")

    if not all([text, source, target]):
        return jsonify({"error": "Missing fields"}), 400

    if MODEL is None:
        try:
            load_model(
                args.model_dir,
                args.base_model_id,
                args.quantize,
            )
        except Exception as e:
            log.exception("Model load failed")
            return jsonify({"error": str(e)}), 500

    try:
        result = translate(text, source, target)
        return jsonify({"translation": result})
    except Exception as e:
        log.exception("Inference failed")
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_dir", required=True)
    parser.add_argument("--base_model_id", default="Qwen/Qwen2.5-0.5B")
    parser.add_argument("--quantize", action="store_true")
    parser.add_argument("--port", type=int, default=8011)
    args = parser.parse_args()

    if not os.path.isdir(args.model_dir):
        log.error("Invalid model directory")
        sys.exit(1)

    log.info(f"Starting Translation API on port {args.port}")
    app.run(host="0.0.0.0", port=args.port, threaded=True)