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

MODEL_ID = "rawsun00001/banking-sms-json-parser-v6-merged"

print("🔄 Loading banking‑SMS JSON parser model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto" if torch.cuda.is_available() else None
)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
print("✅ Model loaded successfully!")

def parse_banking_sms(raw_text: str) -> dict:
    sms_text = " ".join(raw_text.strip().split())
    prompt = sms_text + "|"
    inputs = tokenizer(prompt, return_tensors="pt")
    if torch.cuda.is_available():
        inputs = {k: v.cuda() for k, v in inputs.items()}
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=120,
            do_sample=False,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
    json_part = decoded[len(prompt):].strip()
    m = re.search(r"\{[^{}]+\}", json_part)
    if m:
        try:
            parsed = json.loads(m.group())
            return {
                "date": parsed.get("date"),
                "type": parsed.get("type"),
                "amount": parsed.get("amount"),
                "category": parsed.get("category"),
                "last4": parsed.get("last4"),
                "is_transaction": parsed.get("is_transaction", False),
            }
        except json.JSONDecodeError:
            pass
    return {
        "date": None, "type": None,
        "amount": None, "category": None,
        "last4": None, "is_transaction": False
    }

def predict(raw_text: str) -> str:
    parsed = parse_banking_sms(raw_text)
    if parsed["is_transaction"]:
        summary = (
            f"✅ Transaction Detected!\n\n"
            f"- 📅 Date: **{parsed.get('date', 'N/A')}**\n"
            f"- 💳 Type: **{parsed.get('type', '').title()}**\n"
            f"- 💰 Amount: **{parsed.get('amount')}**\n"
            f"- 🏪 Category: **{parsed.get('category')}**\n"
            f"- 🔢 Last 4 Digits: **{parsed.get('last4')}**\n\n"
            "**Full Parsed JSON:**\n```json\n"
            f"{json.dumps(parsed, indent=2)}\n```"
        )
    else:
        summary = (
            "ℹ️ Non‑transactional SMS / Promotional / Info message.\n\n"
            "**Parsed Classification JSON:**\n```json\n"
            f"{json.dumps(parsed, indent=2)}\n```"
        )
    return summary

iface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(
        lines=3,
        placeholder="Paste your banking SMS or email here…",
        label="Input SMS / Email"
    ),
    outputs=gr.Markdown(label="Parsed Output"),
    title="🏦 Banking SMS JSON Parser",
    description=(
        "Paste any banking SMS (or email) below — the app will detect transaction "
        "vs non-transaction, and display structured JSON output."
    ),
    allow_flagging="never",
    analytics_enabled=False,  # disable Gradio telemetry in Spaces
)

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