from flask import Flask, request, jsonify
import subprocess

MODEL_PATH = "/app/model/qwen2.5-coder-1.5b-instruct-q4_k_m.gguf"
BINARY_PATH = "/llama.cpp/build/bin/llama-cli"

MAX_TOKENS = "128"
THREADS = "2"
CONTEXT_SIZE = "2048"

app = Flask(__name__)

def query_model(system: str, user: str) -> str:
    prompt = f"System: {system}\nUser: {user}" if system else user

    try:
        cmd = [
            BINARY_PATH,
            "-m", MODEL_PATH,
            "-t", THREADS,
            "-c", CONTEXT_SIZE,
            "-p", prompt,
            "--n_predict", MAX_TOKENS,
            "--temp", "0.7",
            "--top_k", "40",
            "--top_p", "0.9",
            "--repeat_penalty", "1.1",
            "--no-warmup"
        ]

        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=90  # allow slightly longer for HF
        )

        if result.returncode != 0:
            return f"Model error:\n{result.stderr.strip()}"

        return result.stdout.strip() or "No output from model."

    except subprocess.TimeoutExpired:
        return "Error: Model took too long (timeout)."
    except Exception as e:
        return f"Unexpected error:\n{str(e)}"


@app.route("/api/query", methods=["GET", "POST"])
def api_query():
    if request.method == "POST":
        data = request.json or {}
        user_prompt = data.get("prompt", "")
        system_prompt = data.get("system", "")
    else:
        user_prompt = request.args.get("prompt", "")
        system_prompt = request.args.get("system", "")

    if not user_prompt.strip():
        return jsonify({"error": "Prompt is required"}), 400

    output = query_model(system_prompt, user_prompt)
    return jsonify({"output": output})


@app.route("/")
def home():
    return "API is running"


@app.route("/health")
def health():
    return {"status": "ok"}


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)