from flask import Flask, request, jsonify
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import os
import json
import time
import uuid

app = Flask(__name__)

# Load config
with open("llm_config.json", "r") as f:
    config = json.load(f).get("openai", {})

REPO_ID = "mradermacher/distilabeled-Hermes-2.5-Mistral-7B-GGUF"
MODEL_FILENAME = "distilabeled-Hermes-2.5-Mistral-7B.Q2_K.gguf"
HF_TOKEN = os.environ.get("HF_API_TOKEN")
CACHE_DIR = "/app/.cache/huggingface"
os.makedirs(CACHE_DIR, exist_ok=True)

MODEL_PATH = hf_hub_download(
    repo_id=REPO_ID,
    filename=MODEL_FILENAME,
    cache_dir=CACHE_DIR,
    token=HF_TOKEN
)

llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=2048,
    n_threads=4,
    n_gpu_layers=0
)

@app.route("/v1/chat/completions", methods=["POST"])
def chat_completions():
    print("Request received at /v1/chat/completions")
    try:
        data = request.get_json(force=True)
        messages = data.get("messages", [])
        max_tokens = int(data.get("max_tokens", config.get("max_tokens", 256)))
        temperature = float(data.get("temperature", config.get("temperature", 0.6)))
        top_p = float(data.get("top_p", config.get("top_p", 0.95)))
        stop = data.get("stop", config.get("stop", ["User:", "Assistant:"]))

        prompt_lines = []
        for msg in messages:
            role = msg.get("role", "").lower()
            content = msg.get("content", "").strip()
            if role == "user":
                prompt_lines.append(f"User: {content}")
            elif role == "assistant":
                prompt_lines.append(f"Assistant: {content}")
        prompt_lines.append("Assistant:")
        prompt = "\n".join(prompt_lines)

        response = llm(
            prompt,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            stop=stop
        )

        result_text = response.get("choices", [{}])[0].get("text", "").strip()

        return jsonify({
            "id": f"chatcmpl-{uuid.uuid4().hex}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": REPO_ID,
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": result_text
                },
                "finish_reason": "stop"
            }]
        })

    except Exception as e:
        return jsonify({"error": str(e)}), 500


if __name__ == "__main__":
    # Run server on all IPs on port 7860 (change if needed)
    app.run(host="0.0.0.0", port=7860)