import gradio as gr
import pandas as pd
from prophet import Prophet
import json
import os
import requests

# Securely load the API key from Hugging Face Secrets
API_KEY = os.environ.get("GOOGLE_API_KEY")
LLM_LOADED = bool(API_KEY)

def create_cases_data():
    """Creates a dummy CSV with customer case data."""
    cases_data = {
        'case_id': [101, 102],
        'customer_name': ['Rohan Gupta', 'Priya Sharma'],
        'customer_query': [
            "The Leather care conditioner 500 ml I ordered (Order #404-7654321-1234567) just arrived, but the bottle is leaking. I want a replacement.",
            "Hi, my order for the Cockpit Cleaner Matt 500 ml (Order #404-1234567-9876543) was supposed to be delivered yesterday, but the tracking says it's delayed. Where is it?"
        ],
        'order_number': ['404-7654321-1234567', '404-1234567-9876543'],
        'item_name': ['Leather care conditioner 500 ml', 'Cockpit Cleaner Matt 500 ml'],
        'order_status': ['Delivered', 'Shipped']
    }
    df_cases = pd.DataFrame(cases_data)
    df_cases.to_csv('customer_cases.csv', index=False)

create_cases_data()

# === USE CASE 1: Future Sales Prediction ===
def predict_future_sales(sku_choice):
    if not sku_choice:
        return "Please select a SKU."
    
    df_sales = pd.read_csv('dummy_sales_history.csv')
    df_sku = df_sales[df_sales['sku'] == sku_choice]

    if df_sku.empty:
        return f"No sales data could be found for '{sku_choice}'. Please check the CSV file."

    df_sku = df_sku.rename(columns={'date': 'ds', 'units_sold': 'y'})
    df_sku['ds'] = pd.to_datetime(df_sku['ds'])

    model_prophet = Prophet(daily_seasonality=False, weekly_seasonality=True, yearly_seasonality=True)
    model_prophet.fit(df_sku)
    future = model_prophet.make_future_dataframe(periods=30)
    forecast = model_prophet.predict(future)

    thirty_day_forecast = forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']].tail(30)
    total_sales = int(thirty_day_forecast['yhat'].sum())
    lower_bound = int(thirty_day_forecast['yhat_lower'].sum())
    upper_bound = int(thirty_day_forecast['yhat_upper'].sum())

    return f"Predicted sales for **{sku_choice}** in the next 30 days: **{total_sales} units**.\nPrediction Range: *{lower_bound} to {upper_bound} units*."

# === USE CASE 2: Automated Case Reply Generation ===
def get_open_cases():
    if not os.path.exists('customer_cases.csv'):
        return []
    df_cases = pd.read_csv('customer_cases.csv')
    return [f"Case {row['case_id']}: {row['customer_name']}" for index, row in df_cases.iterrows()]

def get_case_details(case_selection):
    if not case_selection:
        return "", "", gr.update(visible=False)

    case_id = int(case_selection.split(':')[0].replace('Case', '').strip())
    df_cases = pd.read_csv('customer_cases.csv')
    case_data = df_cases[df_cases['case_id'] == case_id].iloc[0]

    return case_data['customer_query'], case_data.to_json(), gr.update(visible=True)

def generate_case_reply(case_data_json):
    if not LLM_LOADED:
        return "Google API Key not loaded. Please configure it in the Space secrets."

    # --- THIS IS THE CORRECTED URL ---
    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.0-pro:generateContent?key={API_KEY}"
    
    internal_data = json.loads(case_data_json)
    customer_query = internal_data.pop('customer_query')
    internal_data.pop('case_id')

    system_prompt = """
        You are an expert AI assistant for "RecoEngine", specializing in Amazon customer support... (Your prompt remains the same)
        """
    user_prompt = f"### Customer Query ###\n{customer_query}\n### Internal Data ###\n{json.dumps(internal_data)}"

    payload = { "contents": [{"parts": [{"text": f"{system_prompt}\n{user_prompt}"}]}] }
    headers = {'Content-Type': 'application/json'}

    try:
        response = requests.post(url, headers=headers, json=payload)
        response.raise_for_status()
        result = response.json()
        generated_text = result['candidates'][0]['content']['parts'][0]['text']
    except requests.exceptions.RequestException as e:
        return f"Error calling API: {e}"
    except (KeyError, IndexError) as e:
        return f"Error parsing API response: {e}\nRaw Response:\n{result}"

    try:
        start_index = generated_text.find('{')
        end_index = generated_text.rfind('}') + 1
        json_output_str = generated_text[start_index:end_index]
        parsed_json = json.loads(json_output_str)

        template = f"""
          Subject: Regarding your recent inquiry about your order... (Your email template remains the same)
          """
        return template
    except Exception as e:
        return f"Error parsing API output: {e}\nRaw Output:\n{generated_text}"

def send_reply(case_selection):
    case_id = int(case_selection.split(':')[0].replace('Case', '').strip())
    df_cases = pd.read_csv('customer_cases.csv')
    df_cases = df_cases[df_cases['case_id'] != case_id]
    df_cases.to_csv('customer_cases.csv', index=False)
    return "Reply sent, case closed.", "", "", gr.update(choices=get_open_cases(), value=None), gr.update(visible=False)

# === USE CASE 3: Automated Prep Center Communication ===
def generate_prep_center_email(client_name, shipment_id, status, arrival_date, num_cartons, labels_provided):
    if not LLM_LOADED:
        return "Google API Key not loaded. Please configure it in the Space secrets."

    # --- THIS IS THE CORRECTED URL ---
    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.0-pro:generateContent?key={API_KEY}"
    
    shipment_data = {"client_name": client_name, "shipment_id": shipment_id, "status": status, "expected_arrival_date": arrival_date, "num_cartons": num_cartons, "labels_provided": labels_provided}
    system_prompt = "You are an expert AI operations assistant..."
    user_prompt = f"Analyze this data and generate the JSON: {json.dumps(shipment_data)}"
    payload = { "contents": [{"parts": [{"text": f"{system_prompt}\n{user_prompt}"}]}] }
    headers = {'Content-Type': 'application/json'}

    try:
        response = requests.post(url, headers=headers, json=payload)
        response.raise_for_status()
        result = response.json()
        generated_text = result['candidates'][0]['content']['parts'][0]['text']
    except requests.exceptions.RequestException as e:
        return f"Error calling API: {e}"
    except (KeyError, IndexError) as e:
        return f"Error parsing API response: {e}\nRaw Response:\n{result}"
        
    try:
        start_index = generated_text.find('{')
        end_index = generated_text.rfind('}') + 1
        json_output_str = generated_text[start_index:end_index]
        parsed_json = json.loads(json_output_str)
        template = f"""
            Subject: {parsed_json.get('urgency_tag')} - Action Required for Your Shipment... (Your email template remains the same)
            """
        return template
    except Exception as e:
        return f"Error parsing API output: {e}\nRaw Output:\n{generated_text}"

# === USE CASE 4: Best Warehouse Recommendation ===
def recommend_best_warehouse(sku_choice):
    if not sku_choice:
        return "Please select a SKU.", ""

    df_wh = pd.read_csv('daily_sales_history.csv')
    df_wh_sku = df_wh[df_wh['sku'] == sku_choice]

    if df_wh_sku.empty:
        return f"SKU '{sku_choice}' not found in the daily sales data.", ""

    warehouse_sales = df_wh_sku.groupby('warehouse')['quantity_sold'].sum()
    best_warehouse = warehouse_sales.idxmax()
    recommendation = f"The best warehouse for **{sku_choice}** is **{best_warehouse}** with {warehouse_sales.max()} units sold historically."

    df_wh_future = df_wh[(df_wh['sku'] == sku_choice) & (df_wh['warehouse'] == best_warehouse)]
    df_wh_future = df_wh_future.rename(columns={'date': 'ds', 'quantity_sold': 'y'})

    if len(df_wh_future) < 2:
        return recommendation, "Not enough historical data for this specific warehouse to make a prediction."

    model_prophet_wh = Prophet(daily_seasonality=False, weekly_seasonality=True, yearly_seasonality=True)
    model_prophet_wh.fit(df_wh_future)
    future_df = model_prophet_wh.make_future_dataframe(periods=30)
    forecast = model_prophet_wh.predict(future_df)
    future_sales = int(forecast['yhat'].tail(30).sum())

    prediction_text = f"Predicted sales for **{sku_choice}** from **{best_warehouse}** warehouse in the next 30 days: **{future_sales} units**."
    return recommendation, prediction_text

# --- 4. GRADIO UI CONSTRUCTION ---
with gr.Blocks(theme=gr.themes.Soft(text_size='lg'), title="RecoEngine Demo") as demo:
    gr.Markdown("# RecoEngine Demo!!")

    with gr.Tabs():
        # --- TAB 1: Future Sales Prediction ---
        with gr.TabItem("Future Sales Prediction"):
            gr.Markdown("## Predict Next 30-Day Sales for a SKU")
            sku_dropdown_sales = gr.Dropdown(
                choices=['Leather care conditioner 500 ml', 'Cockpit Cleaner Matt 500 ml', 'Cockpit Cleaner Shine 500 ml'],
                label="Select a SKU"
            )
            predict_button_sales = gr.Button("Generate Prediction")
            sales_output = gr.Markdown()
            predict_button_sales.click(predict_future_sales, inputs=sku_dropdown_sales, outputs=sales_output)

        # --- TAB 2: Automated Case Reply ---
        with gr.TabItem("Automated Case Reply"):
            gr.Markdown("## Handle Customer Support Cases")
            case_dropdown = gr.Dropdown(choices=get_open_cases(), label="Select an Open Case")
            with gr.Group():
                gr.Markdown("**Customer's Message:**")
                customer_query_box = gr.Textbox(lines=4, interactive=False, show_label=False)
                case_data_state = gr.State()
            generate_reply_button = gr.Button("Generate AI Reply", visible=False)
            with gr.Group():
                gr.Markdown("**Generated Reply:**")
                reply_output_box = gr.Textbox(lines=8, interactive=False, show_label=False)
            send_reply_button = gr.Button("Send Reply & Close Case")
            status_box = gr.Markdown()
            case_dropdown.change(get_case_details, inputs=case_dropdown, outputs=[customer_query_box, case_data_state, generate_reply_button])
            generate_reply_button.click(generate_case_reply, inputs=case_data_state, outputs=reply_output_box)
            send_reply_button.click(send_reply, inputs=case_dropdown, outputs=[status_box, customer_query_box, reply_output_box, case_dropdown, generate_reply_button])

        # --- TAB 3: Prep Center Communication ---
        with gr.TabItem("Prep Center Communication"):
            gr.Markdown("## Generate Prep Center Reminder Email")
            with gr.Row():
                client_name_input = gr.Textbox(label="Client Name", value="FabFurnish")
                shipment_id_input = gr.Textbox(label="Shipment ID", value="SHP-IND-9001")
            with gr.Row():
                status_input = gr.Dropdown(
                    choices=["In Transit", "Awaiting Arrival", "Arrived", "Processing", "Completed"],
                    label="Status",
                    value="In Transit"
                )
                arrival_date_input = gr.Textbox(label="Expected Arrival Date", value="2025-09-11")
            with gr.Row():
                cartons_input = gr.Number(label="Number of Cartons", value=15)
                labels_provided_input = gr.Checkbox(label="Labels Provided?", value=False)
            generate_email_button = gr.Button("Generate Email")
            email_output = gr.Textbox(lines=10, label="Generated Email")
            generate_email_button.click(generate_prep_center_email, inputs=[client_name_input, shipment_id_input, status_input, arrival_date_input, cartons_input, labels_provided_input], outputs=email_output)

        # --- TAB 4: Best Warehouse Recommendation ---
        with gr.TabItem("Best Warehouse Recommendation"):
            gr.Markdown("## Find the Best Warehouse and Predict Sales")
            sku_dropdown_wh = gr.Dropdown(
                choices=['Leather care conditioner 500 ml', 'Cockpit Cleaner Matt 500 ml', 'Cockpit Cleaner Shine 500 ml'],
                label="Select a SKU"
            )
            recommend_button_wh = gr.Button("Get Recommendation & Prediction")
            wh_recommendation_output = gr.Markdown()
            wh_prediction_output = gr.Markdown()
            recommend_button_wh.click(recommend_best_warehouse, inputs=sku_dropdown_wh, outputs=[wh_recommendation_output, wh_prediction_output])

# --- 5. LAUNCH THE APP ---
if __name__ == "__main__":
    demo.launch()