from flask import Blueprint, request, jsonify
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import os
import json

api = Blueprint("api", __name__)

# Load config from llm_config.json
with open("llm_config.json", "r") as f:
    config = json.load(f)["openai"]

# Model details
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)

# Download model if not already cached
MODEL_PATH = hf_hub_download(
    repo_id=REPO_ID,
    filename=MODEL_FILENAME,
    cache_dir=CACHE_DIR,
    token=HF_TOKEN
)

# Load model
llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=2048,
    n_threads=4,
    n_gpu_layers=0  # Adjust >0 for GPU acceleration if needed
)

@api.route("/v1/chat/completions", methods=["POST"])
def chat_completions():
    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", "")
        content = msg.get("content", "")
        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)

    # Generate completion
    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({
        "choices": [{
            "message": {
                "role": "assistant",
                "content": result_text
            },
            "finish_reason": "stop",
            "index": 0
        }]
    })