from flask import Flask, request, Response, render_template
from pmbl import PMBL
import os
import json

app = Flask(__name__, template_folder='.')
pmbl = PMBL("./PMB-7b.Q6_K.gguf")  # Replace with the path to your model

@app.route("/chat", methods=['POST'])
def chat():
    try:
        data = request.get_json()
        user_input = data["user_input"]
        mode = data["mode"]
        history = pmbl.get_chat_history(mode, user_input)
        response_generator = pmbl.generate_response(user_input, history, mode)
        return Response(response_generator, mimetype='text/plain')
    except Exception as e:
        print(f"[SYSTEM] Error: {str(e)}")
        return json.dumps({"error": str(e)})

@app.route("/", methods=['GET'])
def root():
    return render_template("index.html")

@app.route("/sleep", methods=['POST'])
def sleep():
    try:
        pmbl.sleep_mode()
        return json.dumps({"message": "Sleep mode completed successfully"})
    except Exception as e:
        print(f"[SYSTEM] Error: {str(e)}")
        return json.dumps({"error": str(e)})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 8000)))