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

# === Конфигурация ===
AI_MODEL = "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO"
HF_SPACE_URL = "https://nymbo-serverless-textgen-hub.hf.space"

# === Инициализация ===
app = Flask(__name__)
logging.basicConfig(level=logging.INFO)

try:
    gr_client = Client(src=HF_SPACE_URL)
except Exception as e:
    logging.error(f"Ошибка подключения к HuggingFace Space: {e}")
    gr_client = None

# === Корневой эндпоинт ===
@app.route("/", methods=["GET"])
def home():
    return "Proxy free test", 200

# === Основной OpenAI-подобный endpoint ===
@app.route("/v1/chat/completions", methods=["POST"])
def chat():
    try:
        data = request.get_json()
        messages = data.get("messages", [])
        user_msg = next((m["content"] for m in reversed(messages) if m["role"] == "user"), None)
        system_msg = next((m["content"] for m in messages if m["role"] == "system"), "You are a helpful AI assistant.")
    except Exception as e:
        return jsonify({"error": f"Невалидный JSON: {e}"}), 400

    if not user_msg or not gr_client:
        return jsonify({"error": "Отсутствует сообщение пользователя или не подключена модель."}), 400

    try:
        output = gr_client.predict(
            history=[[user_msg, None]],
            system_msg=system_msg,
            max_tokens=512,
            temperature=0.7,
            top_p=0.95,
            freq_penalty=0,
            seed=-1,
            custom_model=AI_MODEL,
            search_term="",
            selected_model=AI_MODEL,
            api_name="/bot"
        )
    except Exception as e:
        return jsonify({"error": f"Ошибка при вызове модели: {e}"}), 500

    return jsonify({
        "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": AI_MODEL,
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": output
                },
                "finish_reason": "stop"
            }
        ],
        "usage": {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0
        }
    })

# === Запуск сервера ===
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=True)