"""
Gradio web app for the ai-words model.
Runs on HuggingFace Spaces with GPU support.
"""
import os
import subprocess
import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, GenerationConfig
import time
import torch

MODEL_DIR = "./trained_model"
TOKENIZER_NAME = "Qwen/Qwen3-8B"

# Detect device
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

# ─── Load or train model ─────────────────────────────────────────────

def get_model():
    """Load existing trained model, or train from scratch if not available."""
    if os.path.isdir(MODEL_DIR) and os.path.isfile(os.path.join(MODEL_DIR, "config.json")):
        print(f"Loading trained model from {MODEL_DIR}...")
    else:
        print("No trained model found. Training from scratch...")
        # Run the full pipeline (tokenize + train)
        result = subprocess.run(
            ["python", "tokenize_data.py"],
            capture_output=False
        )
        if result.returncode != 0:
            raise RuntimeError("Dataset loading/tokenization failed")

        result = subprocess.run(
            ["python", "train.py"],
            capture_output=False
        )
        if result.returncode != 0:
            raise RuntimeError("Training failed")

    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
    tokenizer.pad_token = tokenizer.eos_token

    generator = pipeline(
        "text-generation",
        model=MODEL_DIR,
        tokenizer=tokenizer,
        device=device,
        clean_up_tokenization_spaces=False,
    )
    return tokenizer, generator


print(f"Using device: {device}")
tokenizer, generator = get_model()
print("Model ready!")


# ─── Generation function ─────────────────────────────────────────────

def generate(prompt, max_tokens, temperature, do_sample):
    if not prompt.strip():
        return "Please enter a prompt.", ""

    config = GenerationConfig(
        max_new_tokens=int(max_tokens),
        do_sample=do_sample,
        temperature=float(temperature) if do_sample else 1.0,
    )

    input_ids = tokenizer.encode(prompt)
    tokens_in = len(input_ids)

    start_time = time.time()
    result = generator(prompt, generation_config=config)
    elapsed = time.time() - start_time

    generated_text = result[0]["generated_text"]
    output_ids = tokenizer.encode(generated_text)
    tokens_out = len(output_ids)
    new_tokens = tokens_out - tokens_in
    speed = new_tokens / elapsed if elapsed > 0 else 0

    stats = (
        f"Prompt tokens: {tokens_in}\n"
        f"New tokens: {new_tokens}\n"
        f"Total output tokens: {tokens_out}\n"
        f"Time: {elapsed:.2f}s\n"
        f"Speed: {speed:.1f} tokens/sec\n"
        f"Device: {device}"
    )

    return generated_text, stats


# ─── Gradio UI ────────────────────────────────────────────────────────

with gr.Blocks(title="ai-words", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🧠 ai-words")
    gr.Markdown(
        "GPT-2 fine-tuned on 13 datasets (~2.9M examples): "
        "WikiText, dictionaries, news, reviews, slang, and more."
    )

    with gr.Row():
        with gr.Column(scale=3):
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="Type your prompt here...",
                lines=3,
                value="The meaning of life is"
            )
            with gr.Row():
                max_tokens = gr.Slider(10, 500, value=100, step=10, label="Max new tokens")
                temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.1, label="Temperature")
                do_sample = gr.Checkbox(value=True, label="Sample")

            generate_btn = gr.Button("Generate", variant="primary")

        with gr.Column(scale=2):
            stats_output = gr.Textbox(label="Stats", lines=6, interactive=False)

    output_text = gr.Textbox(label="Generated Text", lines=8, interactive=False)

    generate_btn.click(
        fn=generate,
        inputs=[prompt_input, max_tokens, temperature, do_sample],
        outputs=[output_text, stats_output],
    )

    gr.Markdown(
        "### Training Data\n"
        "WikiText-103 · WikiText-2 · English Dictionary · WordNet · "
        "AG News · IMDB · Rotten Tomatoes · CNN/DailyMail · Yelp Reviews · "
        "Urban Dictionary · Slang · Gen Z Slang"
    )


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)