import gradio as gr
from transformers import AutoModelForCausalLM
import torch

# Загрузка модели DeepSeek-Coder-1.3b-instruct
model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"
model = AutoModelForCausalLM.from_pretrained(model_name)

# Проверка доступности GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# Функция для генерации кода
def generate_code(prompt, file_type):
    if file_type in ["Gradio", "Vercel", "Streamlit"]:
        # Генерация кода для платформ
        prompt_with_file_type = f"Write a configuration or setup code for {file_type} to: {prompt}"
    else:
        # Генерация кода для языков программирования
        prompt_with_file_type = f"Write a {file_type} code for: {prompt}"
    
    # Генерация кода с использованием модели
    inputs = prompt_with_file_type  # Передаем текст напрямую
    outputs = model.generate(
        inputs,  # Модель должна поддерживать прямой ввод текста
        max_length=200,  # Ограничение длины вывода
        num_return_sequences=1,  # Один вариант ответа
        temperature=0.7,  # Контроль случайности
        top_p=0.9,  # Контроль разнообразия
    )
    
    # Возвращаем сгенерированный код
    return outputs[0]  # Предполагаем, что модель возвращает текст напрямую

# Функция для обновления кода
def update_code(existing_code, update_prompt):
    # Объединение существующего кода и нового запроса
    prompt_with_update = f"Rewrite the following code to: {update_prompt}\n\nExisting Code:\n{existing_code}"
    
    # Генерация обновленного кода
    outputs = model.generate(
        prompt_with_update,
        max_length=250,  # Увеличен для обновления кода
        num_return_sequences=1,
        temperature=0.7,
        top_p=0.9,
    )
    
    # Возвращаем обновленный код
    return outputs[0]

# Примеры для интерфейса
examples = [
    ["Write a function to calculate factorial", "Python"],
    ["Create a simple interface for a calculator", "Gradio"],
    ["Deploy a Next.js app", "Vercel"],
    ["Create a data visualization app", "Streamlit"],
    ["Write a program to reverse a string", "JavaScript"],
    ["Create a responsive navbar", "HTML"],
]

# Gradio интерфейс
with gr.Blocks(theme='Nymbo/Nymbo-Theme') as demo:
    gr.Markdown("# AI Code Generator with Update Feature")
    gr.Markdown("Enter a prompt and select the file type or platform to generate code. You can also update the generated code with a new prompt.")
    
    with gr.Row():
        input_prompt = gr.Textbox(label="Input Prompt", placeholder="e.g., Write a function to calculate factorial...")
        file_type = gr.Dropdown(
            label="File Type / Platform",
            choices=["Python", "JavaScript", "HTML", "CSS", "Java", "C++", "Gradio", "Vercel", "Streamlit"],
            value="Python"
        )
    
    # Панель для кода с очень большим количеством строк
    output_code = gr.Textbox(label="Generated Code", lines=30, interactive=False)
    
    # Кнопка "Скопировать"
    copy_button = gr.Button("Скопировать код")
    copy_button.click(
        None,  # Не требует Python-функции
        inputs=[output_code],  # Входные данные — текст из output_code
        outputs=None,  # Нет выходных данных
        js="""(text) => {
            navigator.clipboard.writeText(text);
            alert('Код скопирован в буфер обмена!');
        }"""
    )
    
    generate_button = gr.Button("Generate Code")
    generate_button.click(fn=generate_code, inputs=[input_prompt, file_type], outputs=output_code)

    # Секция для обновления кода
    with gr.Row():
        update_prompt = gr.Textbox(label="Update Prompt", placeholder="e.g., Add error handling to the code...")
        update_button = gr.Button("Update Code")
    
    update_button.click(fn=update_code, inputs=[output_code, update_prompt], outputs=output_code)

    # Добавление примеров
    gr.Examples(
        examples=examples,
        inputs=[input_prompt, file_type],
        outputs=output_code,
        fn=generate_code,
        cache_examples=True,  # Кэширование для ускорения
        label="Click on an example to get started!"
    )

# Запуск интерфейса
demo.launch()