import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from huggingface_hub import login
import PyPDF2
import pandas as pd
import torch
import os

# Set page configuration
st.set_page_config(
    page_title="WizNerd Insp",
    page_icon="🚀",
    layout="centered"
)

# Load Hugging Face token from environment variable
HF_TOKEN = os.getenv("HF_TOKEN")  # Set this in your Space's secrets

# Model name
MODEL_NAME = "amiguel/instruct_BERT-base-uncased_model"

# Label mapping (same as in Colab)
LABEL_TO_CLASS = {
    0: "Campaign", 1: "Corrosion Monitoring", 2: "Flare Tip", 3: "Flare TIP",
    4: "FU Items", 5: "Intelligent Pigging", 6: "Lifting", 7: "Non Structural Tank",
    8: "Piping", 9: "Pressure Safety Device", 10: "Pressure Vessel (VIE)",
    11: "Pressure Vessel (VII)", 12: "Structure", 13: "Flame Arrestor"
}

# Title with rocket emojis
st.title("🚀 WizNerd Insp 🚀")

# Configure Avatars
USER_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/9904d9a0d445ab0488cf7395cb863cce7621d897/USER_AVATAR.png"
BOT_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/991f4c6e4e1dc7a8e24876ca5aae5228bcdb4dba/Ataliba_Avatar.jpg"

# Sidebar configuration
with st.sidebar:
    st.header("Upload Documents 📂")
    uploaded_file = st.file_uploader(
        "Choose a PDF or XLSX file",
        type=["pdf", "xlsx"],
        label_visibility="collapsed"
    )

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# File processing function
@st.cache_data
def process_file(uploaded_file):
    if uploaded_file is None:
        return ""
    
    try:
        if uploaded_file.type == "application/pdf":
            pdf_reader = PyPDF2.PdfReader(uploaded_file)
            return "\n".join([page.extract_text() for page in pdf_reader.pages])
        elif uploaded_file.type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet":
            df = pd.read_excel(uploaded_file)
            return df.to_markdown()
    except Exception as e:
        st.error(f"📄 Error processing file: {str(e)}")
        return ""

# Model loading function
@st.cache_resource
def load_model(hf_token):
    try:
        if not hf_token:
            st.error("🔐 Authentication required! Please set the HF_TOKEN environment variable.")
            return None
        
        login(token=hf_token)
        
        # Load tokenizer and model for classification
        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=hf_token)
        model = AutoModelForSequenceClassification.from_pretrained(
            MODEL_NAME,
            num_labels=len(LABEL_TO_CLASS),  # Ensure correct number of labels
            token=hf_token
        )
        
        # Determine device
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model.to(device)
        
        return model, tokenizer
        
    except Exception as e:
        st.error(f"🤖 Model loading failed: {str(e)}")
        return None

# Classification function
def classify_instruction(prompt, file_context, model, tokenizer):
    full_prompt = f"Context:\n{file_context}\n\nInstruction: {prompt}"
    
    model.eval()
    device = model.device
    
    inputs = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True, max_length=128)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    with torch.no_grad():
        outputs = model(**inputs)
        prediction = outputs.logits.argmax().item()
        class_name = LABEL_TO_CLASS[prediction]
    
    return class_name

# Display chat messages
for message in st.session_state.messages:
    try:
        avatar = USER_AVATAR if message["role"] == "user" else BOT_AVATAR
        with st.chat_message(message["role"], avatar=avatar):
            st.markdown(message["content"])
    except:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

# Chat input handling
if prompt := st.chat_input("Ask your inspection question..."):
    # Load model if not already loaded
    if "model" not in st.session_state:
        model_data = load_model(HF_TOKEN)
        if model_data is None:
            st.error("Failed to load model. Please ensure HF_TOKEN is set correctly.")
            st.stop()
            
        st.session_state.model, st.session_state.tokenizer = model_data
    
    model = st.session_state.model
    tokenizer = st.session_state.tokenizer
    
    # Add user message
    with st.chat_message("user", avatar=USER_AVATAR):
        st.markdown(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})

    # Process file context
    file_context = process_file(uploaded_file)
    
    # Classify the instruction
    if model and tokenizer:
        try:
            with st.chat_message("assistant", avatar=BOT_AVATAR):
                predicted_class = classify_instruction(prompt, file_context, model, tokenizer)
                response = f"The Item Class is: {predicted_class}"
                st.markdown(response)
                st.session_state.messages.append({"role": "assistant", "content": response})
                
        except Exception as e:
            st.error(f"⚡ Classification error: {str(e)}")
    else:
        st.error("🤖 Model not loaded!")