# app.py
import os
import re
import json
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# ✅ Your model repo (exact)
MODEL_ID = "Amey9766/llama32B-hospitality-review-triage"

# If your model is private/gated, add HF_TOKEN as a Space Secret
HF_TOKEN = os.getenv("HF_TOKEN", None)


def extract_first_json(text: str) -> str:
    """
    Extract the first JSON object from a string. Falls back to raw text.
    """
    match = re.search(r"\{.*\}", text, flags=re.S)
    return match.group(0) if match else text


def load_model():
    """
    Load tokenizer/model once. Keep it simple: no 4-bit, no bitsandbytes.
    This avoids the bitsandbytes / quantization crashes on Spaces.
    """
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)

    # Ensure pad token is set
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        token=HF_TOKEN,
        device_map="auto",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    )
    model.eval()
    return tokenizer, model


tokenizer, model = load_model()


def triage(review_text: str, max_new_tokens: int = 256, temperature: float = 0.0):
    if not review_text or not review_text.strip():
        return "Please enter a hotel review."

    messages = [
        {
            "role": "system",
            "content": (
                "You are a hospitality review triage assistant. "
                "Output ONLY valid JSON (no extra text)."
            ),
        },
        {"role": "user", "content": review_text.strip()},
    ]

    prompt = tokenizer.apply_chat_template(messages, tokenize=False)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    do_sample = float(temperature) > 0.0

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=do_sample,
            temperature=float(temperature) if do_sample else None,
            pad_token_id=tokenizer.eos_token_id,
        )

    decoded = tokenizer.decode(output_ids[0], skip_special_tokens=True)

    # Try to isolate JSON and pretty print it
    json_text = extract_first_json(decoded)
    try:
        obj = json.loads(json_text)
        return json.dumps(obj, indent=2)
    except Exception:
        return json_text


# ✅ Gradio UI
with gr.Blocks(title="Hospitality Review Triage") as demo:
    gr.Markdown("# 🏨 Hospitality Review → JSON Triage")
    gr.Markdown(
        "Paste a guest review and get structured JSON for routing/triage. "
        "The model is instructed to output **JSON only**."
    )

    review_in = gr.Textbox(
        label="Guest Review",
        lines=6,
        placeholder="Example: The room was dirty and the AC didn’t work. Front desk didn’t respond.",
    )

    with gr.Row():
        max_tokens = gr.Slider(64, 512, value=256, step=32, label="Max new tokens")
        temp = gr.Slider(0.0, 1.0, value=0.0, step=0.1, label="Temperature")

    output = gr.Code(label="JSON Output", language="json")

    btn = gr.Button("Generate JSON", variant="primary")
    btn.click(triage, inputs=[review_in, max_tokens, temp], outputs=output)

    gr.Examples(
        examples=[
            ["The room was dirty and the AC didn’t work. I called twice and no one came."],
            ["Great location and staff were friendly, but breakfast was overpriced and slow."],
            ["I found bugs in the bathroom. This is unacceptable and I want a refund."],
            ["Noise from the hallway kept us awake all night. The bed was uncomfortable."],
        ],
        inputs=review_in,
        label="Try examples",
    )

demo.launch()