import gradio as gr
from huggingface_hub import InferenceClient
import os
import fastapi
import uvicorn

# --- Конфигурация ---
# Системный промпт для настройки поведения модели
SYSTEM_PROMPT = """
You are a world-class expert in Linux system administration.
Your task is to convert a user's request in natural language into a SINGLE, executable shell command.
- Do NOT provide any explanation.
- Do NOT use markdown or any other formatting.
- Do NOT suggest multiple options.
- Just return the raw command, ready to be executed.

Example:
User: list all running docker containers
You: docker ps
"""

# Модель для генерации команд (8B версия, не требующая сторонних ключей)
MODEL = "meta-llama/Meta-Llama-3-8B-Instruct"


# --- Логика приложения ---
def generate_command(user_request):
    """
    Обращается к Hugging Face Inference API для генерации команды.
    Токен считывается из секрета HF_TOKEN, заданного в настройках Space.
    """
    # Считываем токен из переменной окружения (секрета Hugging Face)
    hf_token = os.environ.get("HF_TOKEN")

    # Проверки
    if not user_request:
        return "Ошибка: Запрос не может быть пустым."
    if not hf_token:
        return "Ошибка: Секрет HF_TOKEN не найден. Пожалуйста, убедитесь, что он добавлен в настройках репозитория Hugging Face Space."

    try:
        # Выполняем запрос к API
        client = InferenceClient(model=MODEL, token=hf_token)
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_request},
        ]
        response = client.chat_completion(
            messages=messages,
            max_tokens=200,
            temperature=0.1,
        )
        command = response.choices[0].message.content.strip()
        
        # Очистка вывода от артефактов
        if command.startswith("`") and command.endswith("`"):
            command = command[1:-1]
        if command.startswith("bash\n"):
            command = command.replace("bash\n", "", 1)

        return command

    except Exception as e:
        return f"Произошла ошибка API: {str(e)}"


# --- Интерфейс Gradio ---
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🤖 AI System Administrator Assistant
        Введите ваш запрос на естественном языке, и ИИ преобразует его в готовую к выполнению shell-команду.
        **Внимание:** Этот инструмент — ваш помощник. Всегда проверяйте предложенные команды перед выполнением в реальной системе.
        """
    )
    
    # Поле для ввода токена полностью убрано
    user_input = gr.Textbox(label="Ваш запрос", placeholder="Например: 'покажи все файлы в текущей папке, отсортированные по размеру'", lines=3)
    
    generate_button = gr.Button("Сгенерировать команду")
    
    output_command = gr.Textbox(label="Предложенная команда", interactive=False, lines=5)

    generate_button.click(
        fn=generate_command,
        inputs=[user_input], # В инпутах больше нет поля для токена
        outputs=output_command
    )
    
    gr.Markdown(
        """
        ---
        *Модель, используемая для генерации: `meta-llama/Meta-Llama-3-8B-Instruct`.*
        """
    )


# --- Запуск приложения ---
app = fastapi.FastAPI()

@app.get("/health")
def health_check():
    return {"status": "ok"}

app = gr.mount_gradio_app(app, demo, path="/")

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