app.py

import streamlit as st
import streamlit.components.v1 as components
import pandas as pd
import matplotlib.pyplot as plt
import os
from datetime import datetime
from model import predict_delay, get_weather_condition
from utils import validate_inputs, generate_heatmap
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.units import inch
from io import BytesIO
from simple_salesforce import Salesforce
import base64
import logging
import json
import requests

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Streamlit app configuration
st.set_page_config(page_title="Delay 🚀", layout="wide")

# Salesforce connection (using environment variables)
try:
    sf_instance_url = os.environ.get("SF_INSTANCE_URL")
    if not sf_instance_url:
        raise ValueError("SF_INSTANCE_URL environment variable is not set")
    if "lightning.force.com" in sf_instance_url:
        logger.warning("SF_INSTANCE_URL contains lightning.force.com; consider using my.salesforce.com for reliable PDF downloads")
    sf = Salesforce(
        username=os.environ.get("SF_USERNAME"),
        password=os.environ.get("SF_PASSWORD"),
        security_token=os.environ.get("SF_SECURITY_TOKEN"),
        instance_url=sf_instance_url
    )
except Exception as e:
    st.error(f"Failed to connect to Salesforce: {str(e)}")
    logger.error(f"Salesforce connection failed: {str(e)}")
    sf = None

# Weather API configuration
WEATHER_API_KEY = os.environ.get("WEATHER_API_KEY")
WEATHER_API_URL = "http://api.openweathermap.org/data/2.5/forecast"

# Title
st.title("Project Delay Predictor 🚀")

# Task options per phase
task_options = {
    "Planning": ["Define Scope", "Resource Allocation", "Permit Acquisition"],
    "Design": ["Architectural Drafting", "Engineering Analysis", "Design Review"],
    "Construction": ["Foundation Work", "Structural Build", "Utility Installation"]
}

# Initialize session state
if 'phase' not in st.session_state:
    st.session_state.phase = ""
if 'task' not in st.session_state:
    st.session_state.task = ""
if 'weather_data' not in st.session_state:
    st.session_state.weather_data = None

# Function to fetch weather data
def fetch_weather_data(project_location, date):
    if not WEATHER_API_KEY:
        logger.error("WEATHER_API_KEY not set")
        return None, "Weather API key not set. Please provide a valid API key."
    try:
        params = {
            "q": project_location,
            "appid": WEATHER_API_KEY,
            "units": "metric"
        }
        response = requests.get(WEATHER_API_URL, params=params)
        response.raise_for_status()
        data = response.json()
        
        # Find the closest forecast to the specified date
        target_date = datetime.strptime(date, "%Y-%m-%d")
        closest_forecast = None
        min_time_diff = float('inf')
        
        for forecast in data['list']:
            forecast_time = datetime.fromtimestamp(forecast['dt'])
            time_diff = abs((forecast_time - target_date).total_seconds())
            if time_diff < min_time_diff:
                min_time_diff = time_diff
                closest_forecast = forecast
        
        if not closest_forecast:
            return None, "No forecast available for the specified date."
        
        # Map weather conditions to impact score
        weather_main = closest_forecast['weather'][0]['main'].lower()
        if 'clear' in weather_main:
            impact_score = 10
        elif 'clouds' in weather_main:
            impact_score = 30 if closest_forecast['clouds']['all'] < 50 else 50
        elif 'rain' in weather_main:
            impact_score = 70 if closest_forecast['rain'].get('3h', 0) < 2.5 else 85
        elif 'storm' in weather_main or 'thunderstorm' in weather_main:
            impact_score = 90
        else:
            impact_score = 50  # Default for other conditions (e.g., fog, snow)
        
        weather_condition = get_weather_condition(impact_score)
        return {
            "weather_impact_score": impact_score,
            "weather_condition": weather_condition,
            "temperature": closest_forecast['main']['temp'],
            "humidity": closest_forecast['main']['humidity']
        }, None
    except Exception as e:
        logger.error(f"Failed to fetch weather data: {str(e)}")
        return None, f"Failed to fetch weather data for {project_location}: {str(e)}"

# Function to format high_risk_phases with flag and alert
def format_high_risk_phases(high_risk_phases):
    formatted = []
    for phase in high_risk_phases:
        flag = "🚩" if phase['risk'] > 75 else ""
        alert = "[Alert]" if phase['risk'] > 75 else ""
        formatted.append(f"{flag} {phase['phase']}: {phase['task']} (Risk: {phase['risk']:.1f}%) {alert}")
    return formatted

# Function to generate PDF
def generate_pdf(input_data, prediction, heatmap_fig):
    buffer = BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=letter)
    styles = getSampleStyleSheet()
    story = []

    # Title
    story.append(Paragraph("Project Delay Prediction Report", styles['Title']))
    story.append(Spacer(1, 12))

    # Input Data
    story.append(Paragraph("Input Data", styles['Heading2']))
    input_fields = [
        f"Project Name: {input_data['project_name']}",
        f"Phase: {input_data['phase']}",
        f"Task: {input_data['task']}",
        f"Current Progress: {input_data['current_progress']}%",
        f"Task Expected Duration: {input_data['task_expected_duration']} days",
        f"Task Actual Duration: {input_data['task_actual_duration']} days",
        f"Workforce Gap: {input_data['workforce_gap']}%",
        f"Workforce Skill Level: {input_data['workforce_skill_level']}",
        f"Workforce Shift Hours: {input_data['workforce_shift_hours']}",
        f"Weather Impact Score: {input_data['weather_impact_score']}",
        f"Weather Condition: {input_data['weather_condition']}",
        f"Weather Forecast Date: {input_data['weather_forecast_date']}",
        f"Project Location: {input_data['project_location']}"
    ]
    for field in input_fields:
        story.append(Paragraph(field, styles['Normal']))
    story.append(Spacer(1, 12))

    # Prediction Results
    story.append(Paragraph("Prediction Results", styles['Heading2']))
    high_risk_text = "<br/>".join(format_high_risk_phases(prediction['high_risk_phases']))
    prediction_fields = [
        f"Delay Probability: {prediction['delay_probability']:.2f}%",
        f"High Risk Phases:<br/>{high_risk_text}",
        f"AI Insights: {prediction['ai_insights']}",
        f"Weather Condition: {prediction['weather_condition']}"
    ]
    for field in prediction_fields:
        story.append(Paragraph(field, styles['Normal']))
    story.append(Spacer(1, 12))

    # Heatmap
    story.append(Paragraph("Delay Risk Heatmap", styles['Heading2']))
    img_buffer = BytesIO()
    heatmap_fig.savefig(img_buffer, format='png', bbox_inches='tight')
    img_buffer.seek(0)
    story.append(Image(img_buffer, width=6*inch, height=2*inch))
    plt.close(heatmap_fig)

    doc.build(story)
    buffer.seek(0)
    return buffer

# Function to save data to Salesforce, including PDF
def save_to_salesforce(input_data, prediction, pdf_buffer):
    if sf is None:
        return "Salesforce connection not established."
    try:
        # Prepare data for Delay_Predictor__c object
        sf_data = {
            "Project_Name__c": input_data["project_name"],
            "Phase__c": input_data["phase"],
            "Task__c": input_data["task"],
            "Current_Progress__c": input_data["current_progress"],
            "Task_Expected_Duration__c": input_data["task_expected_duration"],
            "Task_Actual_Duration__c": input_data["task_actual_duration"],
            "Workforce_Gap__c": input_data["workforce_gap"],
            "Workforce_Skill_Level__c": input_data["workforce_skill_level"],
            "Workforce_Shift_Hours__c": input_data["workforce_shift_hours"],
            "Weather_Impact_Score__c": input_data["weather_impact_score"],
            "Weather_Condition__c": input_data["weather_condition"],
            "Weather_Forecast_Date__c": input_data["weather_forecast_date"],
            "Project_Location__c": input_data["project_location"],
            "Delay_Probability__c": prediction["delay_probability"],
            "AI_Insights__c": prediction["ai_insights"],
            "High_Risk_Phases__c": "; ".join(format_high_risk_phases(prediction["high_risk_phases"]))
        }
        # Create a new record in Delay_Predictor__c
        result = sf.Delay_Predictor__c.create(sf_data)
        if not result["success"]:
            return f"Salesforce save failed: {result['errors']}"

        # Get the record ID
        record_id = result["id"]

        # Upload PDF as ContentVersion
        pdf_data = pdf_buffer.getvalue()
        pdf_base64 = base64.b64encode(pdf_data).decode('utf-8')
        content_version = {
            "Title": f"Delay_Prediction_Report_{input_data['project_name']}_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
            "PathOnClient": "project_delay_report.pdf",
            "VersionData": pdf_base64,
            "FirstPublishLocationId": record_id
        }
        cv_result = sf.ContentVersion.create(content_version)
        if not cv_result["success"]:
            return f"Failed to upload PDF to Salesforce: {cv_result['errors']}"

        # Get the ContentVersion ID
        content_version_id = cv_result["id"]

        # Query the ContentDocumentId from the ContentVersion
        query = f"SELECT ContentDocumentId FROM ContentVersion WHERE Id = '{content_version_id}'"
        query_result = sf.query(query)
        if query_result["totalSize"] == 0:
            logger.error(f"Failed to retrieve ContentDocumentId for ContentVersion {content_version_id}")
            return "Failed to retrieve ContentDocumentId for the ContentVersion"
        content_document_id = query_result["records"][0]["ContentDocumentId"]

        # Construct the Salesforce URL for the ContentDocument
        pdf_url = f"{sf_instance_url}/sfc/servlet.shepherd/document/download/{content_document_id}"
        logger.info(f"Generated PDF URL: {pdf_url}")

        # Update the Delay_Predictor__c record with the PDF URL
        update_result = sf.Delay_Predictor__c.update(record_id, {"PDF_Report__c": pdf_url})
        if update_result != 204:
            logger.error(f"Failed to update PDF_Report__c with URL: {pdf_url}")
            return f"Failed to update PDF_Report__c field: {update_result}"

        return None
    except Exception as e:
        logger.error(f"Error saving to Salesforce: {str(e)}")
        return f"Error saving to Salesforce: {str(e)}"

# Input form
with st.form("project_form"):
    col1, col2 = st.columns(2)
    
    with col1:
        project_name = st.text_input("Project Name")
        phase = st.selectbox("Phase", [""] + ["Planning", "Design", "Construction"], index=0, key="phase_select")
        
        if phase != st.session_state.phase:
            st.session_state.phase = phase
            st.session_state.task = ""
        task_options_list = [""] + task_options.get(phase, []) if phase else [""]
        task = st.selectbox("Task", task_options_list, index=0, key="task_select")
        current_progress = st.number_input("Current Progress (%)", min_value=0.0, max_value=100.0, step=1.0, value=0.0)
        task_expected_duration = st.number_input("Task Expected Duration (days)", min_value=0, step=1, value=0)
        task_actual_duration = st.number_input("Task Actual Duration (days)", min_value=0, step=1, value=0)
    
    with col2:
        workforce_gap = st.number_input("Workforce Gap (%)", min_value=0.0, max_value=100.0, step=1.0, value=0.0)
        workforce_skill_level = st.selectbox("Workforce Skill Level", ["", "Low", "Medium", "High"], index=0)
        workforce_shift_hours = st.number_input("Workforce Shift Hours", min_value=0, step=1, value=0)
        st.write(f"**Selected Shift Hours**: {workforce_shift_hours}")
        project_location = st.text_input("Project Location (City)", placeholder="e.g., New York")
        weather_forecast_date = st.date_input("Weather Forecast Date", min_value=datetime(2025, 1, 1), value=None)
    
    submit_button = st.form_submit_button("Fetch Weather and Predict Delay")

# Process form submission
if submit_button:
    logger.info("Processing form submission")
    input_data = {
        "project_name": project_name,
        "phase": phase,
        "task": task,
        "current_progress": current_progress,
        "task_expected_duration": task_expected_duration,
        "task_actual_duration": task_actual_duration,
        "workforce_gap": workforce_gap,
        "workforce_skill_level": workforce_skill_level,
        "workforce_shift_hours": workforce_shift_hours,
        "weather_impact_score": 0,  # Placeholder, to be updated
        "weather_condition": "",    # Placeholder, to be updated
        "weather_forecast_date": weather_forecast_date.strftime("%Y-%m-%d") if weather_forecast_date else "",
        "project_location": project_location
    }
    
    # Validate inputs (excluding weather fields initially)
    error = validate_inputs(input_data)
    if error and not error.startswith("Please select or fill in weather"):
        st.error(error)
        logger.error(f"Validation error: {error}")
    else:
        # Fetch weather data
        if project_location and weather_forecast_date:
            weather_data, weather_error = fetch_weather_data(project_location, input_data["weather_forecast_date"])
            if weather_error:
                st.error(weather_error)
                logger.error(weather_error)
                input_data["weather_impact_score"] = 50  # Fallback value
                input_data["weather_condition"] = "Unknown"
            else:
                input_data["weather_impact_score"] = weather_data["weather_impact_score"]
                input_data["weather_condition"] = weather_data["weather_condition"]
                st.write(f"**Weather Data for {project_location} on {input_data['weather_forecast_date']}**:")
                st.write(f"- Condition: {weather_data['weather_condition']}")
                st.write(f"- Impact Score: {weather_data['weather_impact_score']}")
                st.write(f"- Temperature: {weather_data['temperature']}°C")
                st.write(f"- Humidity: {weather_data['humidity']}%")
                st.session_state.weather_data = weather_data
        else:
            st.error("Please provide a project location and weather forecast date.")
            logger.error("Project location or weather forecast date missing")
            input_data["weather_impact_score"] = 50  # Fallback value
            input_data["weather_condition"] = "Unknown"
        
        # Re-validate with weather data
        error = validate_inputs(input_data)
        if error:
            st.error(error)
            logger.error(f"Validation error: {error}")
        else:
            with st.spinner("Generating predictions and AI insights..."):
                try:
                    prediction = predict_delay(input_data)
                except Exception as e:
                    st.error(f"Prediction failed: {str(e)}")
                    logger.error(f"Prediction failed: {str(e)}")
                    prediction = {"error": str(e)}
            
            if "error" in prediction:
                st.error(prediction["error"])
            else:
                st.subheader("Prediction Results")
                st.write(f"**Delay Probability**: {prediction['delay_probability']:.2f}%")
                st.write("**High Risk Phases**:")
                for line in format_high_risk_phases(prediction['high_risk_phases']):
                    st.write(line)
                st.write(f"**AI Insights**: {prediction['ai_insights']}")
                st.write(f"**Weather Condition**: {prediction['weather_condition']}")
                
                # Generate Chart.js heatmap
                chart_config = generate_heatmap(prediction['delay_probability'], f"{phase}: {task}")
                chart_id = f"chart-{hash(str(chart_config))}"
                chart_html = f"""
                <canvas id="{chart_id}" style="max-height: 200px; max-width: 600px;"></canvas>
                <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
                <script>
                    try {{
                        const ctx = document.getElementById('{chart_id}').getContext('2d');
                        new Chart(ctx, {json.dumps(chart_config)});
                    }} catch (e) {{
                        console.error('Chart.js failed: ' + e);
                    }}
                </script>
                """
                try:
                    components.html(chart_html, height=250)
                    logger.info("Chart.js heatmap rendered")
                except Exception as e:
                    logger.error(f"Chart.js rendering failed: {str(e)}")
                    st.error("Failed to render heatmap; please check your browser settings.")
                
                # Generate matplotlib figure for PDF
                fig, ax = plt.subplots(figsize=(8, 2))
                color = 'red' if prediction['delay_probability'] > 75 else 'yellow' if prediction['delay_probability'] > 50 else 'green'
                ax.barh([f"{phase}: {task}"], [prediction['delay_probability']], color=color, edgecolor='black')
                ax.set_xlim(0, 100)
                ax.set_xlabel("Delay Probability (%)")
                ax.set_title("Delay Risk Heatmap")
                plt.tight_layout()
                
                pdf_buffer = generate_pdf(input_data, prediction, fig)
                plt.close(fig)
                st.download_button(
                    label="Download Prediction Report (PDF)",
                    data=pdf_buffer,
                    file_name="project_delay_report.pdf",
                    mime="application/pdf"
                )
                
                # Save to Salesforce, including PDF
                sf_error = save_to_salesforce(input_data, prediction, pdf_buffer)
                if sf_error:
                    st.error(sf_error)
                    logger.error(f"Salesforce error: {sf_error}")
                else:
                    st.success("Prediction data and PDF successfully saved to Salesforce!")
                    logger.info("Data and PDF saved to Salesforce")
                
                st.session_state.prediction = prediction
                st.session_state.input_data = input_data