import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Set page title and icon
st.set_page_config(page_title="SecureFin AI Analyzer", page_icon="🛡️")

# --- MODEL LOADING ---
@st.cache_resource
def load_model():
    model_id = "zoraiz112/SecureFin-SLM-1.5B-Final"
    # Load tokenizer and model
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    # Use device_map="auto" to handle CPU or GPU automatically in the Space
    model = AutoModelForCausalLM.from_pretrained(
        model_id, 
        torch_dtype=torch.float32, # CPU-friendly
        device_map="auto"
    )
    return tokenizer, model

st.title("🛡️ SecureFin AI: Fraud Detection Agent")
st.markdown("Enter transaction details below for a deep-learning security analysis.")

# Load the model (this shows a spinner while loading)
with st.spinner("Initializing SecureFin Engine... (this may take a minute)"):
    tokenizer, model = load_model()

# --- SIDEBAR INPUTS ---
st.sidebar.header("Transaction Details")
amount = st.sidebar.text_input("Amount ($)", "5000.00")
location = st.sidebar.text_input("Location", "Unknown IP / Foreign Country")
category = st.sidebar.selectbox("Category", ["Crypto Exchange", "High-Value Tech", "ATM Withdrawal", "Grocery", "Other"])
time = st.sidebar.text_input("Time of Day", "3:45 AM")

# --- ANALYSIS LOGIC ---
if st.button("Analyze for Fraud"):
    # Create the prompt for the fine-tuned model
    input_text = f"""Analyze this transaction for potential fraud:
    - Amount: ${amount}
    - Location: {location}
    - Category: {category}
    - Time: {time}
    
    Status:"""
    
    with st.spinner("Analyzing patterns..."):
        # Tokenize and Generate
        inputs = tokenizer(input_text, return_tensors="pt")
        # Move to same device as model
        inputs = {k: v.to(model.device) for k, v in inputs.items()}
        
        output_tokens = model.generate(
            **inputs, 
            max_new_tokens=200, 
            temperature=0.1, 
            do_sample=True
        )
        
        response = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
        
        # Display the result
        st.subheader("Analysis Result")
        # Clean up the output to only show the model's new text
        cleaned_response = response.split("Status:")[-1].strip()
        st.info(cleaned_response)

st.divider()
st.caption("SecureFin AI v1.0 | Built on Qwen-2.5-1.5B-Merged")