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

MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"

print("Loading tokenizer...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

print("Loading model...")

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    dtype=torch.float16
)

# IMPORTANT FOR ZEROGPU
model = model.to("cuda")
model.eval()

print("Model loaded successfully!")
print("CUDA available:", torch.cuda.is_available())


@spaces.GPU(duration=120)
def generate_product_content(
    product_name,
    category,
    material,
    color,
    features,
    target_customer
):

    prompt = f"""
You are an ecommerce product copywriter.

Create product content using ONLY the information provided.

Product name: {product_name}
Category: {category}
Material: {material}
Color: {color}
Features: {features}
Target customer: {target_customer}

Return exactly this format:

SHORT_DESCRIPTION:
Write 1-2 concise sentences.

DESCRIPTION:
Write an 80-120 word product description.

KEY_FEATURES:
- Feature 1
- Feature 2
- Feature 3
- Feature 4

SEO_TITLE:
Maximum 60 characters.

META_DESCRIPTION:
Maximum 155 characters.

Rules:
- Do not invent specifications.
- Do not invent dimensions.
- Do not invent certifications.
- Do not make medical claims.
- Do not make unrealistic guarantees.
- Do not mention AI.
- Use natural ecommerce language.
- Use only the information supplied.
"""

    messages = [
        {
            "role": "system",
            "content": "You are a professional ecommerce product copywriter."
        },
        {
            "role": "user",
            "content": prompt
        }
    ]

    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt"
    )

    # Input must also be on CUDA
    inputs = inputs.to("cuda")

    with torch.no_grad():

        outputs = model.generate(
            inputs,
            max_new_tokens=300,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id
        )

    generated_tokens = outputs[0][inputs.shape[-1]:]

    result = tokenizer.decode(
        generated_tokens,
        skip_special_tokens=True
    )

    return result.strip()


with gr.Blocks(title="Product Content AI") as demo:

    gr.Markdown(
        """
        # Product Content AI

        Generate product descriptions and SEO content.
        """
    )

    with gr.Row():

        with gr.Column():

            product_name = gr.Textbox(
                label="Product Name",
                placeholder="Blue Crystal Necklace"
            )

            category = gr.Textbox(
                label="Category",
                placeholder="Necklace"
            )

            material = gr.Textbox(
                label="Material",
                placeholder="Alloy"
            )

            color = gr.Textbox(
                label="Color",
                placeholder="Blue and Gold"
            )

            features = gr.Textbox(
                label="Features",
                placeholder="Crystal pendant, lightweight, adjustable chain",
                lines=4
            )

            target_customer = gr.Textbox(
                label="Target Customer",
                placeholder="Women"
            )

            generate_button = gr.Button(
                "Generate Content",
                variant="primary"
            )

        with gr.Column():

            output = gr.Textbox(
                label="Generated Content",
                lines=18
            )

    generate_button.click(
        fn=generate_product_content,
        inputs=[
            product_name,
            category,
            material,
            color,
            features,
            target_customer
        ],
        outputs=output,
        api_name="generate_product_content"
    )


demo.launch(
    server_name="0.0.0.0",
    server_port=7860
)