from flask import Flask, request, jsonify
import json
import logging
import time
from gradio_client import Client

app = Flask(__name__)

# Configure logging
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Initialize JARVIS client
jarvis = Client("hadadrjt/ai")

# Define available models
models = [
    "JARVIS: 2.1.2",
    "DeepSeek: V3-0324",
    "DeepSeek: R1 (Reasoning)",
    "DeepSeek: R1 - Distill Qwen 14B (Reasoning)",
    "DeepSeek: R1 - Distill Qwen 32B (Reasoning)",
    "DeepSeek: R1 - Distill Llama 70B (Reasoning)",
    "Google: Gemini 2.0 Flash Thinking (Reasoning)",
    "Google: Gemini 2.5 Pro",
    "Google: Gemma 3 1B-IT",
    "Google: Gemma 3 4B-IT",
    "Google: Gemma 3 27B-IT",
    "Meta: Llama 3.1 8B Instruct",
    "Meta: Llama 3.2 3B Instruct",
    "Meta: Llama 3.3 70B Instruct",
    "Meta: Llama 4 Maverick 17B 128E Instruct",
    "Meta: Llama 4 Scout 17B 16E Instruct",
    "Qwen: Qwen2.5 VL 3B Instruct",
    "Qwen: Qwen2.5 VL 32B Instruct",
    "Qwen: Qwen2.5 VL 72B Instruct",
    "Agentica: Deepcoder 14B Preview"
]

@app.route("/v1/chat/completions", methods=["POST"])
def chat_completions():
    data = request.json
    messages = data.get("messages", [])
    model = data.get("model", "JARVIS: 2.1.2")
    stream = data.get("stream", False)

    # Validate messages
    if not isinstance(messages, list):
        return jsonify({"error": {"message": "'messages' must be a list", "code": "invalid_message_list"}}), 400

    for message in messages:
        if not isinstance(message, dict) or 'role' not in message or 'content' not in message:
            return jsonify({"error": {"message": "Each message must have a 'role' and a 'content'", "code": "invalid_message"}}), 400

    # Extract the last message content
    last_message = messages[-1]["content"]

    try:
        # Get full response with context
        jarvis.predict(new=model, api_name="/change_model")
        result = jarvis.predict(multi={"text": last_message}, api_name="/api")
        response_text = result[0][0][1]

        # Create OpenAI compatible response
        response = {
            "id": f"chatcmpl-{int(time.time())}",
            "object": "chat.completion",
            "created": int(time.time()),
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": response_text
                },
                "finish_reason": "stop"
            }],
            "usage": {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0
            }
        }

        return jsonify(response)
    except Exception as e:
        logger.error(f"Error processing request: {str(e)}")
        return jsonify({"error": {"message": str(e), "code": "jarvis_error"}}), 500

@app.route("/v1/models", methods=["GET"])
def list_models():
    return jsonify({"data": [{"id": model} for model in models], "object": "list"})

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