import streamlit as st
import pandas as pd
from transformers import pipeline

# Load a reasoning LLM
llm = pipeline("text2text-generation", model="google/flan-t5-large")

SUPPLIER_FILE = "supplier_master.xlsx"
CUSTOMER_FILE = "customer_master.xlsx"

def fetch_supplier_master():
    return pd.read_excel(SUPPLIER_FILE)

def fetch_contractor_master():
    return pd.read_excel(CUSTOMER_FILE)

def structured_retrieval(query, supplier_df, contractor_df):
    """Map queries to structured filters instead of raw string search."""
    query_lower = query.lower()
    context = ""

    # Example: suppliers with MSA
    if "msa" in query_lower and "expired" in query_lower:
        expired = supplier_df[supplier_df['Duration'].str.contains("Dec") == False]
        if not expired.empty:
            context += "Suppliers with expired MSA:\n" + expired.to_string(index=False) + "\n\n"
    elif "msa" in query_lower:
        active = supplier_df[supplier_df['MSA Exists'].str.lower() == "yes"]
        if not active.empty:
            context += "Suppliers with MSA:\n" + active.to_string(index=False) + "\n\n"

    # Example: contractor performance
    if "top performer" in query_lower:
        top = contractor_df[contractor_df['Performance Category'].str.lower() == "high"]
        if not top.empty:
            context += "Top performers:\n" + top.to_string(index=False) + "\n\n"
    if "below par" in query_lower or "low performer" in query_lower:
        low = contractor_df[contractor_df['Performance Category'].str.lower() == "low"]
        if not low.empty:
            context += "Below par performers:\n" + low.to_string(index=False) + "\n\n"

    return context if context else "No relevant rows found."

def answer_query(query, context):
    """LLM generates a consolidated answer based on retrieved context."""
    prompt = f"Leadership question: {query}\n\nRelevant data:\n{context}\n\nProvide a clear consolidated answer for leadership."
    result = llm(prompt, max_length=256, do_sample=False)
    return result[0]['generated_text']

def main():
    st.title("Supplier & Contractor Agent")

    supplier_data = fetch_supplier_master()
    contractor_data = fetch_contractor_master()

    st.subheader("Leadership Query Interface")
    query = st.text_input("Enter your query about suppliers or contractors:")

    if query:
        st.info(f"Processing query: {query}")
        context = structured_retrieval(query, supplier_data, contractor_data)
        response = answer_query(query, context)
        st.success("Consolidated Response:")
        st.write(response)

    st.subheader("Quick Performance Snapshot")
    role_distribution = contractor_data.groupby('Role').size()
    budget_utilization = contractor_data['Budget Approved'].str.replace('$','').str.replace('M','').astype(float).sum()

    st.write("📊 Role Distribution", role_distribution)
    st.write(f"💰 Overall Budget Utilization: ${budget_utilization}M")

if __name__ == "__main__":
    main()