import pandas as pd
import gradio as gr
from langchain.agents import initialize_agent, Tool
from langchain_google_genai import ChatGoogleGenerativeAI  # requires GOOGLE_API_KEY set as env var

# Store the uploaded dataframe in memory
dataframes = {}

# -----------------------------
# Tools (Excel Q&A Functions)
# -----------------------------
def get_otb_revenue(month: str) -> str:
    if "df" not in dataframes:
        return "Please upload a file first."
    df = dataframes["df"]
    try:
        otb = df.loc[df['Month'].str.lower() == month.lower(), 'OTB Revenue'].values[0]
        stly = df.loc[df['Month'].str.lower() == month.lower(), 'STLY Revenue'].values[0]
        return f"OTB Revenue for {month}: {otb}, STLY Revenue: {stly}"
    except Exception:
        return f"Could not find OTB/ STLY Revenue for {month}"

def check_occupancy(month: str) -> str:
    if "df" not in dataframes:
        return "Please upload a file first."
    df = dataframes["df"]
    try:
        occ = df.loc[df['Month'].str.lower() == month.lower(), 'Occupancy %'].values[0]
        target = df.loc[df['Month'].str.lower() == month.lower(), 'Budgeted Occupancy %'].values[0]
        if occ >= target:
            return f"Occupancy for {month} is {occ}% vs target {target}% ✅ On track!"
        else:
            return f"Occupancy for {month} is {occ}% vs target {target}% ⚠️ Might not meet target."
    except Exception:
        return f"Could not find occupancy data for {month}"

# -----------------------------
# Upload Excel
# -----------------------------
def upload_excel(file):
    if file is None:
        return "Please upload an Excel file."
    dfs = pd.read_excel(file.name, sheet_name=None)
    dataframes.clear()
    dataframes.update(dfs)
    sheet_info = []
    for name, df in dfs.items():
        sheet_info.append(f"'{name}': {len(df)} rows, {len(df.columns)} cols")
    return f"✅ Loaded sheets: {', '.join(sheet_info)}"

# -----------------------------
# LangChain Agent Setup
# -----------------------------
tools = [
    Tool(
        name="Get OTB Revenue",
        func=get_otb_revenue,
        description="Get OTB Revenue and STLY Revenue for a given month (e.g. 'August')"
    ),
    Tool(
        name="Check Occupancy",
        func=check_occupancy,
        description="Check occupancy vs target for a given month (e.g. 'August')"
    )
]

llm = ChatGoogleGenerativeAI(model="gemini-2.5-pro", temperature=0)  # Requires GOOGLE_API_KEY in environment
agent = initialize_agent(tools, llm, agent="zero-shot-react-description", verbose=True)

def chat_agent(message, history):
    if "df" not in dataframes:
        return "Please upload a file first."
    try:
        return agent.run(message)
    except Exception as e:
        return f"⚠️ Agent error: {e}"

# -----------------------------
# Gradio UI
# -----------------------------
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 Excel Agent (LangChain + Gemini)")
    with gr.Row():
        with gr.Column():
            file_input = gr.File(label="Upload Excel file", type="filepath")
            upload_output = gr.Textbox(label="Upload status")
        file_input.upload(upload_excel, file_input, upload_output)

    gr.ChatInterface(fn=chat_agent, title="Excel Chat Agent")

demo.launch()