import logging
from flask import Flask, request, jsonify, render_template
from flask_cors import CORS
from gradio_client import Client
import os
import traceback

app = Flask(__name__)
CORS(app)  # Enable CORS for all routes
logging.basicConfig(level=logging.DEBUG)

# Initialize the client (this loads the API)
CLIENT_URL = os.getenv("CLIENT_URL")
MODEL = os.getenv("MODEL")

if not CLIENT_URL:
    raise ValueError("CLIENT_URL environment variable must be set")
if not MODEL:
    raise ValueError("MODEL environment variable must be set")

app.logger.info(f"Initializing client with URL: {CLIENT_URL}")
client = Client('Nymbo/Groq-Playground-Master')

# System message
system_message = os.getenv("SYS")

context = [{"role": "system", "content": system_message}]

def chat_with_ai(message):
    app.logger.debug(f"Received message: {message}")
    context.append({"role": "user", "content": message})

    # Prepare the full prompt
    full_prompt = "\n".join([f"{msg['role']}: {msg['content']}" for msg in context])

    app.logger.debug(f"Making API call with prompt: {full_prompt[:100]}...")  # Log first 100 chars of prompt
    try:
        result = client.predict(
            message=full_prompt,
            request=MODEL,
            param_3=0.5,
            param_4=8192,
            param_5=0.5,
            param_6=0,
            api_name="/chat"
        )
        app.logger.debug(f"Received result: {result[:100]}...")  # Log first 100 chars of result
    except Exception as e:
        app.logger.error(f"Error in API call: {str(e)}")
        app.logger.error(traceback.format_exc())
        raise

    # Add AI response to context
    context.append({"role": "assistant", "content": result})

    return result

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/chat', methods=['POST'])
def chat():
    data = request.json
    user_input = data.get('user_input')
    
    if not user_input:
        return jsonify({"error": "No user input provided"}), 400
    
    app.logger.debug(f"Received request: {data}")
    
    try:
        response = chat_with_ai(user_input)
        app.logger.debug(f"API response content: {response[:100]}...")  # Log first 100 chars of response
        return jsonify({"response": response})
    except Exception as e:
        app.logger.error(f"Unexpected error: {str(e)}")
        app.logger.error(traceback.format_exc())
        return jsonify({"error": f"An unexpected error occurred: {str(e)}"}), 500

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