import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr
import json
import os
from datetime import datetime
import firebase_admin
from firebase_admin import credentials, firestore

# === Firebase Setup ===
firebase_key = json.loads(os.environ["FIREBASE_CREDENTIALS_JSON"])
cred = credentials.Certificate(firebase_key)
firebase_admin.initialize_app(cred)
db = firestore.client()

# === Load Model ===
model_checkpoint = "chikki2004/incident-bart-model"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def parse_gradio_result(response_text):
    attributes = {}
    if ";" in response_text:
        lines = response_text.strip().split(";")
    else:
        lines = response_text.strip().split("\n")
    for line in lines:
        if ":" in line:
            key, value = line.split(":", 1)
            attributes[key.strip().lower().replace(" ", "_")] = value.strip()
    return attributes

def predict(input_text):
    # Predict
    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding="max_length", max_length=256)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    outputs = model.generate(**inputs, max_length=256)
    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Log to Firebase
    parsed = parse_gradio_result(decoded)
    db.collection("gradio_predictions").add({
        "input_text": input_text,
        "raw_prediction": decoded,
        "parsed_attributes": parsed,
        "timestamp": datetime.now().isoformat()
    })

    return decoded

# ✅ This line exposes it at /run/predict
iface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=5),
    outputs="text",
    title="Incident Attribute Predictor",
    api_name="/predict"
)

iface.launch(ssr_mode=False,share=True)