from llama_cpp import Llama
import gradio as gr
import os

# Локальный путь, куда будет скачана модель при запуске
MODEL_PATH = "unsloth.Q4_K_M.gguf"

# Если модели ещё нет — скачиваем
if not os.path.exists(MODEL_PATH):
    import requests

    url = "https://huggingface.co/eliza555beth2002/DeepSeek-R1-Distill-Text2SQL-OneEpoch-GGUF-q4/resolve/main/unsloth.Q4_K_M.gguf"
    print("📥 Скачиваю модель из Hugging Face...")
    r = requests.get(url, stream=True)
    with open(MODEL_PATH, "wb") as f:
        for chunk in r.iter_content(chunk_size=8192):
            f.write(chunk)
    print("Модель скачана успешно!")

# Загружаем модель Llama.cpp
llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=4096,
    n_threads=4,
    n_gpu_layers=0,  # если на CPU — ставим 0
    verbose=False,
)

# Основная функция
def text2sql(question):
    schema = """
    create table Authors(
        AuthorID integer not null primary key check(AuthorID>0),
        AuthorFIO varchar(40) not null
    );
    create table Books(
        Cipher integer not null primary key check(Cipher>0),
        BookName varchar(4000) not null,
        BookTheme varchar(30) not null
        check(BookTheme in('Любовь','Дружба','Смерть','Общественные проблемы','Внутренние противоречия')),
        BookGenre varchar(15) not null
        check(BookGenre in('Роман','Поэма','Рассказ','Пьеса','Эпопея','Драма'))
    );
    create table Wrote(
        IDAuthor integer not null check(IDAuthor>0),
        BookCipher integer not null check(BookCipher>0),
        foreign key(IDAuthor) references Authors(AuthorID),
        foreign key(BookCipher) references Books(Cipher),
        primary key(IDAuthor, BookCipher)
    );
    """

    prompt = f"""Ты — помощник, который преобразует естественный язык в SQL.
Используй приведённую схему таблиц.
Вопрос: {question}
Контекст базы данных:
{schema}

Ответь ТОЛЬКО SQL-запросом в формате:
```sql
SELECT ...
```"""

    # Генерация
    output = llm(
        prompt,
        max_tokens=512,
        temperature=0.2,
        stop=["```"],
    )

    text = output["choices"][0]["text"].strip()
    return text

# Интерфейс Gradio
demo = gr.Interface(
    fn=text2sql,
    inputs=gr.Textbox(label="Введите вопрос на естественном языке"),
    outputs=gr.Textbox(label="Сгенерированный SQL-запрос"),
    title="🧠 DeepSeek Text2SQL Demo",
    description="Модель DeepSeek-R1-Distill преобразует текст в SQL-запрос по заданной схеме."
)

if __name__ == "__main__":
    demo.launch()