from flask import Flask, request, jsonify
from hugchat import hugchat
from hugchat.login import Login
import os

app = Flask(__name__)

print("Hello welcome to Sema AI", flush=True)  # Flush to ensure immediate output

# Get Hugging Face credentials from environment variables
email = os.getenv('HF_EMAIL')
password = os.getenv('HF_PASS')

#email = "trianatha3@gmail.com"
#password = "@Tri@n@th@1"
print(f"email is {email} and password is {password}", flush=True)

@app.route("/")
def hello():
    return "hello 🤗, Welcome to Sema AI Chat Service."
    
# Flask route to handle incoming chat requests
@app.route('/chat', methods=['POST'])
def chat():
    # Get JSON data from the POST request
    data = request.json
    prompt = data.get('prompt')
    email = data.get('email')
    password = data.get('password')

    print(f"email 2 is {email} and password 2 is {password} and The user wants to Know: {prompt}", flush=True)
    
    if not (password):
        return jsonify({"error": "Missing password"}), 400
    elif not (prompt):
        return jsonify({"error": "Missing prompt"}), 400
    elif not (email):
        return jsonify({"error": "Missing email"}), 400
    else:
        return jsonify({"error": "Missing prompt, email, or password"}), 400

    # Generate the response
    response = generate_response(prompt, email, password)
    
    # Return the response as JSON
    return jsonify({"response": response})

# Function for generating LLM response
def generate_response(prompt_input, email, passwd):
    # Hugging Face Login
    sign = Login(email, passwd)
    cookies = sign.login()
    # Create ChatBot                        
    chatbot = hugchat.ChatBot(cookies=cookies.get_dict())

    # Simple dialogue structure
    string_dialogue = "You are a helpful assistant."
    string_dialogue += f"\n\nUser: {prompt_input}\n\nAssistant: "

    # Generate and return the response
    return chatbot.chat(string_dialogue)

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