import streamlit as st
import streamlit.components.v1 as components
import pandas as pd
import matplotlib.pyplot as plt
import plotly.figure_factory as ff
import os
from datetime import datetime, timedelta
import json
import requests
import base64
import logging
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

# 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 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, {"error": "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 target 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, {"error": "No forecast available for the specified date."}
        
        # Map weather conditions to impact score
        weather_main = forecast['weather'][0]['main'].lower()
        impact_score = 50  # Default
        if 'clear' in weather_main:
            impact_score = 10
        elif 'clouds' in weather_main:
            impact_score = 30 if forecast['clouds']['all'] < 50 else 50
        elif 'rain' in weather_main:
            impact_score = 70 if forecast['rain'].get('3h', 0) < 2.5 else 85
        elif 'storm' in weather_main or 'thunderstorm' in weather_main:
            impact_score = 90
        
        weather_condition = get_weather_condition(impact_score)
        return {
            "weather_impact_score": impact_score,
            "weather_condition": weather_condition,
            "temperature": forecast['main']['temp'],
            "humidity": forecast['main']['humidity']
        }, None
    except Exception as e:
        logger.error(f"Failed to fetch weather data: {str(e)}")
        return None, {"error": 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 Gantt chart
def generate_gantt_chart(input_data, prediction):
    try:
        phase = input_data["phase"]
        task = input_data["task"]
        expected_duration = input_data["task_expected_duration"]
        actual_duration = input_data["task_actual_duration"]
        forecast_date = datetime.strptime(input_data["weather_forecast_date"], "%Y-%m-%d")
        delay_risk = prediction["delay_probability"]

        # Calculate start and end dates
        start_date = forecast_date - timedelta(days=max(expected_duration, actual_duration))
        expected_end = start_date + timedelta(days=expected_duration)
        actual_end = start_date + timedelta(days=actual_duration) if actual_duration > 0 else expected_end

        # Prepare Gantt chart data
        df = [
            dict(Task=f"{phase}: {task} (Expected)", Start=start_date.strftime("%Y-%m-%d"), Finish=expected_end.strftime("%Y-%m-%d"), Resource="Expected", Risk=delay_risk),
            dict(Task=f"{phase}: {task} (Actual)", Start=start_date.strftime("%Y-%m-%d"), Finish=actual_end.strftime("%Y-%m-%d"), Resource="Actual", Risk=delay_risk)
        ]

        # Color based on delay risk
        colors = {
            "Expected": "rgb(0, 255, 0)" if delay_risk <= 50 else "rgb(255, 255, 0)" if delay_risk <= 75 else "rgb(255, 0, 0)",
            "Actual": "rgb(0, 200, 0)" if delay_risk <= 50 else "rgb(200, 200, 0)" if delay_risk <= 75 else "rgb(200, 0, 0)"
        }

        # Create Gantt chart
        fig = ff.create_gantt(
            df,
            colors=colors,
            index_col="Resource",
            title=f"Gantt Chart for {phase}: {task}",
            show_colorbar=True,
            bar_width=0.4,
            showgrid_x=True,
            showgrid_y=True
        )
        fig.update_layout(
            xaxis_title="Timeline",
            yaxis_title="Task",
            height=300,
            margin=dict(l=150)
        )
        return fig
    except Exception as e:
        logger.error(f"Failed to generate Gantt chart: {str(e)}")
        return None

# Function to generate PDF
def generate_pdf(input_data, prediction, heatmap_fig, gantt_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']))
    
    # Check for 2-week risk alert in AI insights
    two_week_alert = next((insight for insight in prediction['ai_insights'].split("; ") if "2-Week Risk Alert" in insight), None)
    if two_week_alert:
        story.append(Paragraph("2-Week Risk Alert", styles['Heading3']))
        story.append(Paragraph(two_week_alert, styles['Normal']))
        story.append(Spacer(1, 12))
    
    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))
    story.append(Spacer(1, 12))

    # Gantt Chart
    if gantt_fig:
        story.append(Paragraph("Gantt Chart", styles['Heading2']))
        gantt_buffer = BytesIO()
        try:
            gantt_fig.write_image(gantt_buffer, format='PNG')
            gantt_buffer.seek(0)
            story.append(Image(gantt_buffer, width=6*inch, height=3*inch))
        except Exception as e:
            logger.error(f"Failed to include Gantt chart in PDF: {str(e)}")
            story.append(Paragraph("Gantt Chart unavailable due to rendering issues.", styles['Normal']))
        story.append(Spacer(1, 12))

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

# Function to save data to Salesforce, including PDF and Status__c
def save_to_salesforce(input_data, prediction, pdf_buffer):
    if sf is None:
        return "Salesforce connection not established."
    try:
        # Determine Status__c based on delay probability
        status = "Flagged" if prediction["delay_probability"] > 75 else "Running"

        # 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"])),
            "Status__c": status
        }
        logger.info(f"Attempting to save to Salesforce Delay_Predictor__c: {sf_data}")
        
        # Create a new record in Delay_Predictor__c
        result = sf.Delay_Predictor__c.create(sf_data)
        if not result["success"]:
            logger.error(f"Salesforce save failed: {result['errors']}")
            return f"Salesforce save failed: {result['errors']}"

        # Get the record ID
        record_id = result["id"]
        logger.info(f"Created Salesforce record ID: {record_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"]:
            logger.error(f"Failed to upload PDF to Salesforce: {cv_result['errors']}")
            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 section
st.markdown("### Project Details")
col1, col2 = st.columns([1, 1])  # Equal width columns for better alignment

with col1:
    project_name = st.text_input("Project Name", help="Enter the name of the project")
    phase = st.selectbox(
        "Phase",
        [""] + ["Planning", "Design", "Construction"],
        index=0 if st.session_state.phase == "" else ["", "Planning", "Design", "Construction"].index(st.session_state.phase),
        key="phase_select",
        help="Select the project phase"
    )
    
    # Update task options when phase changes
    if phase != st.session_state.phase:
        st.session_state.phase = phase
        st.session_state.task = ""  # Reset task when phase changes
        logger.info(f"Phase changed to {phase}, resetting task")
    
    task_options_list = [""] + task_options.get(phase, []) if phase else [""]
    logger.info(f"Task options for phase '{phase}': {task_options_list}")
    task = st.selectbox(
        "Task",
        task_options_list,
        index=0 if st.session_state.task == "" else task_options_list.index(st.session_state.task) if st.session_state.task in task_options_list else 0,
        key="task_select",
        help="Select the task corresponding to the phase"
    )
    st.session_state.task = task
    current_progress = st.number_input("Current Progress (%)", min_value=0.0, max_value=100.0, step=1.0, value=0.0, help="Enter the current progress percentage")
    task_expected_duration = st.number_input("Task Expected Duration (days)", min_value=0, step=1, value=0, help="Enter the expected duration in days")
    task_actual_duration = st.number_input("Task Actual Duration (days)", min_value=0, step=1, value=0, help="Enter the actual duration in days")

with col2:
    workforce_gap = st.number_input("Workforce Gap (%)", min_value=0.0, max_value=100.0, step=1.0, value=0.0, help="Enter the workforce gap percentage")
    workforce_skill_level = st.selectbox("Workforce Skill Level", ["", "Low", "Medium", "High"], index=0, help="Select the workforce skill level")
    workforce_shift_hours = st.number_input("Workforce Shift Hours", min_value=0, step=1, value=0, help="Enter the shift hours")
    st.write(f"**Selected Shift Hours**: {workforce_shift_hours}")
    project_location = st.text_input("Project Location (City)", placeholder="e.g., New York", help="Enter the city for weather data")
    weather_forecast_date = st.date_input("Weather Forecast Date", min_value=datetime(2025, 1, 1), value=None, help="Select the forecast date")

# Predict button
predict_button = st.button("Fetch Weather and Predict Delay")

# Process inputs when button is clicked
if predict_button:
    logger.info("Processing prediction request")
    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.get("error", "Unknown weather error"))
                logger.error(weather_error.get("error", "Unknown 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()
                
                # Generate Gantt chart
                gantt_fig = generate_gantt_chart(input_data, prediction)
                if gantt_fig:
                    st.plotly_chart(gantt_fig, use_container_width=True)
                    logger.info("Gantt chart rendered")
                
                pdf_buffer = generate_pdf(input_data, prediction, fig, gantt_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