import gradio as gr
import json
import os
from datasets import Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
import torch
import tempfile

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

def load_and_train():
    global model, tokenizer
    
    try:
        # Читаем данные напрямую из файла
        with open('training_data.json', 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # Берем первые 100 примеров для быстрого обучения
        sample_data = data[:100]
        
        # Создаем датасет напрямую из данных
        texts = [item['text'] for item in sample_data]
        dataset = Dataset.from_dict({'text': texts})
        
        # Загружаем модель
        model_name = "Salesforce/codegen-350M-mono"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(model_name)
        
        # Устанавливаем pad_token
        tokenizer.pad_token = tokenizer.eos_token
        
        # Подготавливаем данные
        def preprocess_function(examples):
            tokenized = tokenizer(examples["text"], truncation=True, max_length=256, padding="max_length")
            tokenized["labels"] = tokenized["input_ids"].copy()
            return tokenized
        
        processed_dataset = dataset.map(preprocess_function, batched=True, remove_columns=dataset.column_names)
        
        # Параметры обучения
        training_args = TrainingArguments(
            output_dir="./model_output",
            per_device_train_batch_size=1,
            num_train_epochs=2,
            save_steps=20,
            logging_steps=5,
            learning_rate=5e-5,
            max_steps=30,
            report_to="none",
            remove_unused_columns=False,
        )
        
        # Обучаем
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=processed_dataset,
        )
        
        result = trainer.train()
        
        # Сохраняем модель в разрешенную директорию
        model.save_pretrained("./trained_model")
        tokenizer.save_pretrained("./trained_model")
        
        return f"✅ Модель обучена!\nДанных: {len(sample_data)} примеров\nФинальная потеря: {result.training_loss:.4f}\nМодель сохранена в ./trained_model"
        
    except Exception as e:
        return f"❌ Ошибка: {str(e)}"

def generate_code(prompt):
    global model, tokenizer
    
    if model is None:
        return "❌ Модель не обучена. Нажмите 'Обучить модель' сначала."
    
    try:
        # Формируем промпт
        full_prompt = f"World of Warcraft addon: {prompt}\n\nCode:\n"
        
        # Токенизация
        inputs = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True, max_length=256)
        
        # Генерация
        with torch.no_grad():
            outputs = model.generate(
                inputs["input_ids"],
                max_length=400,
                num_return_sequences=1,
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        # Декодирование
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Извлекаем код
        code = generated_text.split("Code:\n")[-1].strip()
        
        return code
        
    except Exception as e:
        return f"❌ Ошибка генерации: {str(e)}"

def check_data():
    try:
        with open('training_data.json', 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        return f"✅ Данные доступны\nВсего записей: {len(data)}\nПример:\n{data[0]['text'][:200]}..."
        
    except Exception as e:
        return f"❌ Ошибка: {str(e)}"

# Создаем интерфейс
with gr.Blocks() as demo:
    gr.Markdown("# 🎮 WoW Addon AI Generator")
    
    with gr.Tabs():
        with gr.TabItem("Обучение"):
            gr.Markdown("### Обучение модели на ваших данных")
            
            with gr.Row():
                check_btn = gr.Button("Проверить данные")
                train_btn = gr.Button("Обучить модель", variant="primary")
            
            train_status = gr.Textbox(label="Статус обучения", lines=8)
            
            check_btn.click(fn=check_data, outputs=train_status)
            train_btn.click(fn=load_and_train, outputs=train_status)
        
        with gr.TabItem("Генерация кода"):
            gr.Markdown("### Генерация кода аддонов")
            
            with gr.Row():
                with gr.Column():
                    prompt_input = gr.Textbox(
                        label="Опишите аддон:",
                        placeholder="Например: Аддон для отслеживания здоровья цели",
                        lines=3
                    )
                    generate_btn = gr.Button("Сгенерировать код", variant="primary")
                
                with gr.Column():
                    code_output = gr.Code(
                        label="Сгенерированный код (Lua):",
                        language="lua",
                        lines=15
                    )
            
            # Примеры
            gr.Examples(
                examples=[
                    "Аддон для отслеживания здоровья цели",
                    "Аддон для подсчета урона в секунду",
                    "Аддон для отображения координат",
                    "Аддон для управления инвентарем"
                ],
                inputs=prompt_input
            )
            
            generate_btn.click(fn=generate_code, inputs=prompt_input, outputs=code_output)

demo.launch()