import json
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import spaces

# Configuration
MODEL_NAME = "speakleash/Bielik-11B-v2.3-Instruct"
# DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DEVICE = "cuda"
TORCH_DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
MAX_TOKENS = 1000

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=TORCH_DTYPE).to(
    DEVICE
)

# Load prompts
with open("prompts.json") as f:
    prompts = json.load(f)


@spaces.GPU
def transform_text(prompt_name, user_input):
    """Transform text using selected prompt and Bielik model"""
    try:
        # Get selected prompt
        selected_prompt = next(p for p in prompts if p["name"] == prompt_name)

        # Create messages structure
        messages = [
            {"role": "system", "content": selected_prompt["system_message"]},
            {"role": "user", "content": user_input},
        ]

        # Tokenize and generate
        input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(
            DEVICE
        )

        generated_ids = model.generate(
            input_ids, max_new_tokens=MAX_TOKENS, do_sample=True
        )

        return tokenizer.batch_decode(generated_ids)[0]

    except Exception as e:
        return f"Error: {str(e)}"


# Create Gradio interface
with gr.Blocks(title="Bielik Goblin") as interface:
    gr.Markdown("# Bielik Goblin")

    with gr.Row():
        prompt_select = gr.Dropdown(
            choices=[p["name"] for p in prompts],
            label="Wybierz prompt",
            interactive=True,
        )

    user_input = gr.Textbox(
        label="Twój tekst", placeholder="Wpisz tutaj swój tekst...", lines=5
    )

    transform_btn = gr.Button("Przekształć tekst", variant="primary")

    with gr.Column():
        progress = gr.StatusTracker()
        output = gr.Textbox(label="Wynik", interactive=False)

    transform_btn.click(
        fn=transform_text,
        inputs=[prompt_select, user_input],
        outputs=output,
        status=progress,
    )

interface.queue().launch(debug=True)