# app.py (get-best-answer Space) — FIXED

import os
from flask import Flask, request, jsonify
from huggingface_hub import InferenceClient

app = Flask(__name__)
MODEL_ID = os.getenv("MODEL_ID", "mistralai/Mixtral-8x7B-Instruct-v0.1")
API_KEY = os.getenv("API_KEY")  # set in Space secrets
client = InferenceClient(token=API_KEY)

def load_instruction():
    with open("instruction_template.txt", "r", encoding="utf-8") as f:
        return f.read()

def generate_summary(context, question, reply_type="brief", language="en",
                     max_new_tokens=800, temperature=0.5):
    lang_map = {"hi":"Hindi", "hindi":"Hindi", "en":"English", "english":"English"}
    target_language = lang_map.get((language or "en").lower(), "English")
    reply_type_instruction = "concise" if reply_type == "brief" else "detailed"

    tmpl = load_instruction()
    prompt = tmpl.format(
        target_language=target_language,
        reply_type_instruction=reply_type_instruction,
        context=context or "",
        question=question or "",
    )

    # Use chat_completion (conversational task)
    completion = client.chat_completion(
        model=MODEL_ID,
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt},
        ],
        max_tokens=max_new_tokens,   # chat uses max_tokens
        temperature=temperature,
        top_p=0.9,
        stop=["</s>", "[/INST]"],
        stream=False,
    )

    text = completion.choices[0].message["content"] if completion and completion.choices else ""
    return (text or "").strip()



@app.post("/summarize")
def summarize():
    data = request.get_json() or {}
    context = data.get("context", "")
    question = data.get("question", "")
    language = data.get("language", "en")
    reply_type = data.get("reply_type", "brief")

    if not question:
        return jsonify({"error": "Missing 'question'"}), 400

    try:
        answer = generate_summary(context, question, reply_type, language)
        if not answer:
            return jsonify({"answer": "I’m momentarily unable to summarize. Please try again."}), 200
        return jsonify({"answer": answer}), 200
    except Exception as e:
        # prevent 500s leaking to your main backend
        return jsonify({"answer": f"I hit an upstream issue: {e}"}), 200

@app.get("/")
def root():
    return {"ok": True, "service": "get-best-answer"}

@app.get("/healthz")
def healthz():
    return {"ok": True}

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