"""
Gradio Space для интерактивного использования модели Code Analyzer
Модель анализирует код студента на основе условия задачи
Поддерживает как веб-интерфейс, так и API endpoint
"""
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import json
from pydantic import BaseModel
from typing import Optional

# Загрузка модели (будет выполнена при первом запуске)
MODEL_NAME = "Vilyam888/Code_analyze.1.0"

def load_model():
    """Загружает модель один раз при старте и сохраняет в глобальные переменные"""
    global model, tokenizer
    
    if model is None or tokenizer is None:
        print("Загрузка модели...")
        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
            device_map="auto",
            trust_remote_code=True
        )
        print("Модель загружена успешно!")

# Глобальные переменные для модели
model, tokenizer = None, None

def build_input(task, code):
    """Форматирует вход в том же формате, что использовался при обучении"""
    parts = []
    if task.strip():
        parts.append(f"Задача:\n{task.strip()}")
    if code.strip():
        parts.append(f"Решение (код):\n```python\n{code.strip()}\n```")
    return "\n\n".join(parts)

def analyze_code(task, code):
    """Анализирует код студента и возвращает результат в формате JSON"""
    global model, tokenizer
    
    if model is None or tokenizer is None:
        return "⏳ Загрузка модели... Пожалуйста, подождите."
    
    if not task.strip():
        return "❌ Пожалуйста, введите условие задачи."
    
    if not code.strip():
        return "❌ Пожалуйста, введите код решения студента."
    
    try:
        # Форматирование входа в стиле обучения
        input_text = build_input(task, code)
        prompt = f"{input_text}\n\nОтвет:\n"
        
        # Токенизация
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        # Генерация ответа
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=600,  # Оптимизировано для CPU
                temperature=0.7,
                top_p=0.85,
                top_k=20,
                repetition_penalty=1.05,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id
            )
        
        # Декодирование ответа
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Извлечение JSON ответа
        if "Ответ:" in response:
            json_str = response.split("Ответ:")[-1].strip()
        else:
            json_str = response
        
        # Попытка распарсить JSON
        try:
            result = json.loads(json_str)
            # Форматируем JSON для красивого отображения
            return json.dumps(result, ensure_ascii=False, indent=2)
        except json.JSONDecodeError:
            # Если не JSON, возвращаем как есть
            return json_str
        
    except Exception as e:
        return f"❌ Ошибка: {str(e)}"

# Создание Gradio интерфейса
with gr.Blocks(title="Code Analyzer Model") as demo:
    gr.Markdown("""
    # 🔍 Code Analyzer Model
    
    Модель для анализа кода студента на основе условия задачи.
    
    **Как использовать:**
    1. Введите условие задачи в первое поле
    2. Введите код решения студента во второе поле
    3. Нажмите "Анализировать код"
    4. Получите детальный анализ в формате JSON
    
    **API Endpoint:** `/api/predict` - используйте для интеграции в другие проекты
    """)
    
    with gr.Row():
        with gr.Column():
            task_input = gr.Textbox(
                label="Условие задачи",
                placeholder="Введите условие задачи...",
                lines=5
            )
            code_input = gr.Textbox(
                label="Код решения студента",
                placeholder="Введите код решения...",
                lines=10
            )
            submit_btn = gr.Button("Анализировать код", variant="primary", size="lg")
        
        with gr.Column():
            result_output = gr.JSON(
                label="Результат анализа"
            )
            # Также показываем как текст для удобства копирования
            result_text = gr.Textbox(
                label="Результат (JSON текст)",
                lines=15,
                interactive=False
            )
    
    # Примеры
    gr.Examples(
        examples=[
            [
                "Напишите функцию, которая принимает список чисел и возвращает сумму всех элементов.",
                "def sum_list(numbers):\n    total = 0\n    for num in numbers:\n        total += num\n    return total"
            ],
            [
                "Создайте функцию для вычисления факториала числа.",
                "def factorial(n):\n    if n <= 1:\n        return 1\n    return n * factorial(n-1)"
            ],
        ],
        inputs=[task_input, code_input]
    )
    
    def analyze_and_format(task, code):
        """Анализирует код и форматирует результат"""
        result_str = analyze_code(task, code)
        try:
            result_json = json.loads(result_str)
            return result_json, result_str
        except:
            return None, result_str
    
    # Загрузка модели при старте
    demo.load(load_model, outputs=None)
    
    # Обработчик кнопки
    submit_btn.click(
        fn=analyze_and_format,
        inputs=[task_input, code_input],
        outputs=[result_output, result_text]
    )
    
    # Обработчик Enter в поле кода
    code_input.submit(
        fn=analyze_and_format,
        inputs=[task_input, code_input],
        outputs=[result_output, result_text]
    )

# Gradio автоматически создает API endpoint на /api/predict
# Для использования из других проектов:
# POST https://YOUR_SPACE_URL.hf.space/api/predict
# Body: {"data": ["task", "code"]}

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(), server_name="0.0.0.0")