import os
import shutil
from flask import Flask, render_template, request, jsonify
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

app = Flask(__name__)

# Kill any process using port 7860
def kill_port(port):
    for proc in psutil.process_iter(attrs=["pid", "connections"]):
        for conn in proc.info["connections"]:
            if conn.laddr.port == port:
                os.kill(proc.info["pid"], 9)

kill_port(7860)  # Ensure Flask doesn't crash due to a used port

# Define cache directory
os.environ["HF_HOME"] = "/app/cache"

# Load Myanmarsar-GPT (1.42B params) from Hugging Face
MODEL_NAME = "simbolo-ai/Myanmarsar-GPT"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=cache_dir)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, cache_dir=cache_dir)

# Function to generate chatbot responses
def generate_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        output = model.generate(**inputs, max_length=200)
    return tokenizer.decode(output[0], skip_special_tokens=True)

# Serve the HTML page
@app.route("/")
def home():
    return render_template("index.html")

# API route for chatbot responses
@app.route("/chat", methods=["POST"])
def chat():
    user_message = request.json.get("message", "")
    bot_reply = generate_response(user_message)
    return jsonify({"reply": bot_reply})

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))  # Default to 7860, but use any assigned port
    app.run(host="0.0.0.0", port=port)