from flask import Flask, render_template, request, jsonify, session
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from datetime import datetime

app = Flask(__name__)
app.secret_key = "supersecretkey"  # Dùng để quản lý session

# Tải mô hình và tokenizer từ Hugging Face
model_path = "phamhai/Llama-3.2-3B-Instruct-Frog"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path)

@app.route('/')
def index():
    session.setdefault('history', [])  # Tạo session nếu chưa có
    return render_template('index.html', history=session['history'])

@app.route('/chat', methods=['POST'])
def chat():
    user_input = request.json.get("message")
    if not user_input:
        return jsonify({"response": "Xin hãy nhập tin nhắn!"})

    # Tokenize input và sinh output
    messages = [
        {"role": "system", "content": "Bạn là trợ lý của tôi, tên là Vivi."},
        {"role": "user", "content": user_input}
    ]
    tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
    outputs = model.generate(tokenized_chat, max_new_tokens=128)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Lưu lại vào lịch sử chat
    chat_entry = {
        "user": user_input,
        "bot": response,
        "timestamp": datetime.now().strftime("%H:%M:%S")  # Thêm thời gian
    }
    session['history'].append(chat_entry)
    session.modified = True  # Cập nhật session

    return jsonify({"response": response, "timestamp": chat_entry["timestamp"]})

@app.route('/clear', methods=['POST'])
def clear():
    session['history'] = []  # Xóa lịch sử chat
    return jsonify({"status": "ok"})

if __name__ == "__main__":
    app.run(debug=True)