# app.py
import gradio as gr
from datetime import datetime
import pytz
import os
from dotenv import load_dotenv
from simple_salesforce import Salesforce
from risk_model import predict_risk
import pandas as pd
import matplotlib.pyplot as plt

# Load environment variables
load_dotenv()

# In-memory prediction history
history = []

# Convert IST to UTC for Salesforce
def format_timestamp_ist_to_utc(timestamp_ist):
    ist = pytz.timezone("Asia/Kolkata")
    dt_ist = datetime.strptime(timestamp_ist, "%Y-%m-%d %H:%M:%S")
    dt_ist = ist.localize(dt_ist)
    dt_utc = dt_ist.astimezone(pytz.utc)
    return dt_utc.strftime("%Y-%m-%dT%H:%M:%SZ")

# Insert prediction into Salesforce
def insert_to_salesforce(temp, duration, risk, timestamp_ist):
    try:
        sf = Salesforce(
            username=os.getenv('SF_USERNAME'),
            password=os.getenv('SF_PASSWORD'),
            security_token=os.getenv('SF_SECURITY_TOKEN'),
            domain='login'
        )
        formatted_timestamp = format_timestamp_ist_to_utc(timestamp_ist)
        sf.Heating_Mantle_Log__c.create({
            'Max_Temperature__c': temp,
            'Duration_Minutes__c': duration,
            'Risk_Level__c': risk,
            'Risk_Timestamp__c': formatted_timestamp
        })
        print("✅ Salesforce record created.")
    except Exception as e:
        print(f"❌ Salesforce error: {e}")

# Main classification logic
def classify(temp, duration):
    if temp <= 0 or duration <= 0:
        return "❌ Invalid", "", "", "", pd.DataFrame(), plt.figure(), ""

    risk_level, risk_score, alert_msg = predict_risk(temp, duration)
    timestamp_ist = datetime.now(pytz.timezone("Asia/Kolkata")).strftime("%Y-%m-%d %H:%M:%S")

    # Log in Salesforce
    insert_to_salesforce(temp, duration, risk_level, timestamp_ist)

    # Save to history
    history.append({
        "Temperature": temp,
        "Duration": duration,
        "Risk": risk_level,
        "Score": risk_score,
        "Timestamp": timestamp_ist
    })

    df = pd.DataFrame(history)

    # Risk trend plot
    risk_map = {'Low': 1, 'Moderate': 2, 'High': 3}
    df["Risk_Num"] = df["Risk"].map(risk_map)
    fig, ax = plt.subplots(figsize=(6, 3))
    ax.plot(df["Timestamp"], df["Risk_Num"], marker="o", linestyle="-", color="crimson")
    ax.set_ylim(0.5, 3.5)
    ax.set_yticks([1, 2, 3])
    ax.set_yticklabels(['Low', 'Moderate', 'High'])
    ax.set_title("Risk Trend Over Time")
    ax.set_xlabel("Timestamp")
    ax.set_ylabel("Risk Level")
    ax.tick_params(axis='x', rotation=45)
    plt.tight_layout()

    df_display = df[["Temperature", "Duration", "Risk", "Score", "Timestamp"]]

    summary_md = f"""
### 🔎 Prediction Summary  
- **Temperature**: {temp} °C  
- **Duration**: {duration} min  
- **Risk Level**: {risk_level}  
- **Risk Score**: {risk_score}%  
- **Alert**: {alert_msg}  
- **Timestamp**: {timestamp_ist} IST  
    """

    return risk_level, f"{risk_score}%", alert_msg, timestamp_ist + " IST", df_display, fig, summary_md

# Reset function
def reset():
    return 100, 30, "", "", "", "", pd.DataFrame(), plt.figure(), ""

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🔥 Heating Mantle Safety Risk Predictor")

    reset_btn = gr.Button("🔄 Reset Inputs")

    with gr.Row():
        temp = gr.Number(label="Max Temperature (°C)", value=100)
        duration = gr.Number(label="Duration (min)", value=30)

    predict_btn = gr.Button("🔍 Predict")

    with gr.Row():
        risk_out = gr.Textbox(label="Risk Level")
        score_out = gr.Textbox(label="Risk Score")
        alert_out = gr.Textbox(label="🚨 Alert Message")
        time_out = gr.Textbox(label="Timestamp (IST)")

    df_display = gr.Dataframe(headers=["Temperature", "Duration", "Risk", "Score", "Timestamp"],
                              label="📈 Prediction History")

    chart = gr.Plot(label="📊 Risk Trend Chart")

    summary = gr.Markdown()

    # Button bindings
    predict_btn.click(fn=classify,
                      inputs=[temp, duration],
                      outputs=[risk_out, score_out, alert_out, time_out, df_display, chart, summary])

    reset_btn.click(fn=reset,
                    outputs=[temp, duration, risk_out, score_out, alert_out, time_out, df_display, chart, summary])

demo.launch()