import os
import json
from flask import Flask, render_template, request, jsonify
from flask_cors import CORS
import requests
from dotenv import load_dotenv

# Загружаем переменные окружения
load_dotenv()

app = Flask(__name__)
CORS(app)

# Проверяем наличие API ключа
API_KEY = os.getenv("MISTRAL_API_KEY")
API_URL = os.getenv("MISTRAL_API_URL", "https://api.mistral.ai/v1/")

# Функция для генерации ответа через API
def generate_response_api(prompt, max_length=1024):
    if not API_KEY:
        return "Ошибка: API ключ не найден. Пожалуйста, добавьте MISTRAL_API_KEY в файл .env"
    
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }
    
    data = {
        "model": "mistral-medium",  # или другая доступная модель
        "messages": [
            {"role": "user", "content": prompt}
        ],
        "max_tokens": max_length,
        "temperature": 0.7,
        "top_p": 0.9
    }
    
    try:
        response = requests.post(f"{API_URL}chat/completions", headers=headers, json=data)
        response.raise_for_status()
        result = response.json()
        return result["choices"][0]["message"]["content"]
    except Exception as e:
        return f"Ошибка при обращении к API: {str(e)}"

# Маршруты
@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/chat', methods=['POST'])
def chat():
    data = request.json
    prompt = data.get('prompt', '')
    
    if not prompt:
        return jsonify({"error": "Пустой запрос"}), 400
    
    try:
        response = generate_response_api(prompt)
        return jsonify({"response": response})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/api/code', methods=['POST'])
def code():
    data = request.json
    prompt = data.get('prompt', '')
    language = data.get('language', 'python')
    
    if not prompt:
        return jsonify({"error": "Пустой запрос"}), 400
    
    # Добавляем контекст для генерации кода
    code_prompt = f"Напиши код на языке {language} для решения следующей задачи: {prompt}"
    
    try:
        response = generate_response_api(code_prompt)
        return jsonify({"code": response})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    print("🚀 Запуск AI Assistant в режиме API")
    print("⚠️ Убедитесь, что вы добавили MISTRAL_API_KEY в файл .env")
    app.run(debug=True, port=5000)