import gradio as gr
import pandas as pd
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
import uvicorn
from sklearn.linear_model import LinearRegression
import base64
import os
from datetime import datetime
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
from reportlab.lib.utils import simpleSplit
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from io import BytesIO
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

weather_map = {"Cloudy": 0, "Rainy": 1, "Sunny": 2}

print("Loading and preprocessing data...")
try:
    if not os.path.exists("new_delay_data.csv"):
        print("Warning: new_delay_data.csv not found. Using default dataset.")
        default_data = {
            "Phase": ["Framing", "Foundation", "Finishing"],
            "Weather": ["Sunny", "Rainy", "Cloudy"],
            "Absentee": [10, 20, 5],
            "DelayLog": [5, 10, 2],
            "Delay%": [30, 60, 15]
        }
        df = pd.DataFrame(default_data)
    else:
        df = pd.read_csv("new_delay_data.csv")

    df = pd.get_dummies(df, columns=["Phase"], drop_first=True)
    df["Weather"] = df["Weather"].map(weather_map)
    df.dropna(subset=["Weather", "Absentee", "DelayLog", "Delay%"], inplace=True)

    for col in ["Phase_Framing", "Phase_Foundation"]:
        if col not in df.columns:
            df[col] = 0
    print("Data loaded and preprocessed. Columns:", df.columns.tolist())
except Exception as e:
    print(f"Error loading data: {e}")
    raise

try:
    X = df[["Phase_Framing", "Phase_Foundation", "Weather", "Absentee", "DelayLog"]]
    y = df["Delay%"]
except Exception as e:
    print(f"Error preparing features: {e}")
    raise

print("Training model...")
try:
    model = LinearRegression()
    model.fit(X, y)
    print("Model trained successfully.")
except Exception as e:
    print(f"Error training model: {e}")
    raise

# New fast AI insight generator without ML model, just logic + template

def generate_ai_insight(phase, weather, absentee_pct, delay_log, prediction):
    # Build a dynamic insight string
    insight = f"Predicted delay of {prediction}% indicates a "
    if prediction >= 75:
        insight += "high risk of project delay. Immediate action is recommended.\n"
    elif prediction >= 50:
        insight += "moderate risk of project delay. Monitor and manage resources carefully.\n"
    else:
        insight += "low risk of project delay. Continue with current plans but remain vigilant.\n"

    insight += f"Phase: {phase}. Weather conditions are {weather.lower()}.\n"

    if absentee_pct > 30:
        insight += f"High absenteeism ({absentee_pct}%) may severely impact progress. Consider temporary staffing or overtime.\n"
    elif absentee_pct > 10:
        insight += f"Moderate absenteeism ({absentee_pct}%) requires attention to maintain productivity.\n"
    else:
        insight += f"Low absenteeism ({absentee_pct}%) supports steady work progress.\n"

    if delay_log > 5:
        insight += f"Previous delays logged ({delay_log}) suggest bottlenecks; analyze root causes and improve workflows.\n"
    else:
        insight += f"Past delays ({delay_log}) are minimal; maintain efficient task coordination.\n"

    if weather == "Rainy":
        insight += "Rainy weather may cause disruptions; plan indoor or protected activities.\n"
    elif weather == "Cloudy":
        insight += "Cloudy weather is less disruptive but monitor conditions.\n"
    else:
        insight += "Sunny weather is ideal for outdoor tasks; maximize on-site work.\n"

    insight += "\n**Suggested Migration Plan:**\n"
    if prediction >= 75:
        insight += "- Deploy additional workforce immediately.\n"
        insight += "- Reschedule non-critical tasks.\n"
        insight += "- Increase monitoring frequency and daily progress reporting.\n"
    elif prediction >= 50:
        insight += "- Cross-train staff to cover absenteeism.\n"
        insight += "- Review supply chains for potential delays.\n"
    else:
        insight += "- Maintain current schedule.\n"
        insight += "- Conduct routine check-ins to prevent issues.\n"

    return insight.strip()

def generate_heatmap(phase, weather, model):
    try:
        absentee_range = np.linspace(0, 100, 20)
        delay_log_range = np.linspace(0, 20, 20)
        framing = 1 if phase == "Framing" else 0
        foundation = 1 if phase == "Foundation" else 0
        weather_encoded = weather_map.get(weather, 0)

        Z = np.zeros((len(delay_log_range), len(absentee_range)))
        for i, delay_log in enumerate(delay_log_range):
            for j, absentee in enumerate(absentee_range):
                input_data = [[framing, foundation, weather_encoded, absentee, delay_log]]
                Z[i, j] = model.predict(input_data)[0]

        plt.figure(figsize=(8, 6))
        sns.heatmap(Z, xticklabels=np.round(absentee_range, 1), yticklabels=np.round(delay_log_range, 1),
                    cmap="YlOrRd", annot=True, fmt=".1f", cbar_kws={'label': 'Predicted Delay %'})
        plt.xlabel("Absentee %")
        plt.ylabel("Previous Delay Log")
        plt.title(f"Delay Prediction Heatmap (Phase: {phase}, Weather: {weather})")

        output_dir = "pdf_reports"
        os.makedirs(output_dir, exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        heatmap_path = os.path.join(output_dir, f"heatmap_{timestamp}.png")
        plt.savefig(heatmap_path, bbox_inches='tight')
        plt.close()
        return heatmap_path
    except Exception as e:
        print(f"Heatmap generation failed: {e}")
        return None

def generate_pdf_report(phase, weather, absentee_pct, delay_log, prediction, risk, insight):
    try:
        buffer = BytesIO()
        c = canvas.Canvas(buffer, pagesize=letter)

        try:
            pdfmetrics.registerFont(TTFont('DejaVuSans', 'DejaVuSans.ttf'))
            c.setFont("DejaVuSans", 12)
            flag_indicator = " 🚩" if prediction >= 75 else ""
        except Exception as e:
            print(f"Failed to load DejaVuSans font: {e}. Falling back to Helvetica with text flag.")
            c.setFont("Helvetica", 12)
            flag_indicator = " [FLAG]" if prediction >= 75 else ""

        c.drawString(100, 750, "Project Delay Prediction Report")
        c.drawString(100, 730, f"Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

        y_position = 700
        max_width = 400

        details = [
            f"Phase: {phase}",
            f"Weather: {weather}",
            f"Absentee Percentage: {absentee_pct}%",
            f"Previous Delay Log: {delay_log}",
            f"Predicted Delay: {prediction}%{flag_indicator}",
            f"Risk Level: {risk}",
            "AI Insight & Migration Plan:"
        ]

        for line in details:
            lines = simpleSplit(line, 'DejaVuSans' if 'DejaVuSans' in pdfmetrics.getRegisteredFontNames() else 'Helvetica', 12, max_width)
            for wrapped_line in lines:
                c.drawString(100, y_position, wrapped_line)
                y_position -= 16

        insight_lines = simpleSplit(insight, 'DejaVuSans' if 'DejaVuSans' in pdfmetrics.getRegisteredFontNames() else 'Helvetica', 12, max_width)
        for wrapped_line in insight_lines:
            c.drawString(100, y_position, wrapped_line)
            y_position -= 16

        heatmap_path = generate_heatmap(phase, weather, model)
        if heatmap_path and os.path.exists(heatmap_path):
            c.drawString(100, y_position - 20, "Delay Prediction Heatmap:")
            c.drawImage(heatmap_path, 100, y_position - 250, width=400, height=200)
            y_position -= 270

        c.showPage()
        c.save()

        pdf_data = buffer.getvalue()
        buffer.close()
        pdf_base64 = base64.b64encode(pdf_data).decode("utf-8")

        output_dir = "pdf_reports"
        os.makedirs(output_dir, exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_path = os.path.join(output_dir, f"delay_report_{timestamp}.pdf")
        with open(output_path, "wb") as f:
            f.write(pdf_data)

        return pdf_base64, output_path, heatmap_path
    except Exception as e:
        print(f"PDF generation failed: {e}")
        return None, None, None

def predict_delay(phase, weather, absentee_pct, delay_log):
    try:
        valid_phases = ["Framing", "Foundation", "Finishing"]
        valid_weather = ["Sunny", "Rainy", "Cloudy"]
        phase = phase if isinstance(phase, str) and phase in valid_phases else "Framing"
        weather = weather if isinstance(weather, str) and weather in valid_weather else "Sunny"
        absentee_pct = float(absentee_pct) if isinstance(absentee_pct, (int, float, str)) and float(absentee_pct) >= 0 else 0
        delay_log = float(delay_log) if isinstance(delay_log, (int, float, str)) and float(delay_log) >= 0 else 0

        framing = 1 if phase == "Framing" else 0
        foundation = 1 if phase == "Foundation" else 0
        weather_encoded = weather_map.get(weather, 0)
        input_data = [[framing, foundation, weather_encoded, absentee_pct, delay_log]]

        prediction = model.predict(input_data)[0]
        prediction = round(prediction, 2)

        if prediction >= 75:
            risk = "High Risk"
        elif prediction >= 50:
            risk = "Moderate Risk"
        else:
            risk = "Low Risk"

        insight = generate_ai_insight(phase, weather, absentee_pct, delay_log, prediction)

        pdf_base64, pdf_path, heatmap_path = generate_pdf_report(phase, weather, absentee_pct, delay_log, prediction, risk, insight)

        return prediction, risk, insight, pdf_base64, pdf_path, heatmap_path
    except Exception as e:
        print(f"Prediction error: {e}")
        return None, None, f"Error: {e}", None, None, None

api_app = FastAPI()

@api_app.post("/predict")
async def predict_from_salesforce(request: Request):
    try:
        data = await request.json()
        phase = data.get("phase", "Framing")
        weather = data.get("weather", "Sunny")
        absentee_pct = data.get("absentee_pct", 0)
        delay_log = data.get("delay_log", 0)

        prediction, risk, insight, pdf_base64, pdf_path, heatmap_path = predict_delay(phase, weather, absentee_pct, delay_log)

        if prediction is None:
            return JSONResponse(status_code=500, content={"status": "error", "message": insight})

        return JSONResponse(content={
            "delay_probability": prediction,
            "risk_alert": risk,
            "ai_insight": insight,
            "pdf_report_base64": pdf_base64 if pdf_base64 else "",
            "pdf_local_path": pdf_path if pdf_path else "PDF generation failed",
            "heatmap_path": heatmap_path if heatmap_path else "Heatmap generation failed",
            "status": "success"
        })
    except Exception as e:
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})

try:
    with gr.Blocks() as demo:
        gr.Markdown("## 🏗️ Delay Predictor with AI Insights (Fast, No ML Text Generation)")
        with gr.Row():
            phase_input = gr.Textbox(label="Phase (Framing/Foundation/Finishing)", value="Framing")
            weather_input = gr.Textbox(label="Weather (Sunny/Rainy/Cloudy)", value="Sunny")
        with gr.Row():
            absentee_input = gr.Number(label="Absentee %", value=0)
            delay_input = gr.Number(label="Previous Delay Log", value=0)
        output = gr.Textbox(label="Prediction Summary")
        submit = gr.Button("Predict")

        def predict_and_format(phase, weather, absentee, delay_log):
            prediction, risk, insight, pdf_base64, pdf_path, heatmap_path = predict_delay(phase, weather, absentee, delay_log)
            if prediction is None:
                return f"Error: {insight}"
            flag = " 🚩" if prediction >= 75 else ""
            return (f"Predicted Delay: {prediction}%{flag}\n"
                    f"Risk Level: {risk}\n"
                    f"Insight & Migration Plan:\n{insight}\n\n"
                    f"PDF Report: {'Saved locally at ' + pdf_path if pdf_path else 'Failed to generate'}\n"
                    f"Heatmap: {'Saved locally at ' + heatmap_path if heatmap_path else 'Failed to generate'}\n"
                    f"PDF Base64: {'Generated' if pdf_base64 else 'Not generated'}")

        submit.click(
            predict_and_format,
            inputs=[phase_input, weather_input, absentee_input, delay_input],
            outputs=output
        )
except Exception as e:
    print(f"Error setting up Gradio UI: {e}")
    raise

try:
    app = gr.mount_gradio_app(api_app, demo, path="/")
except Exception as e:
    print(f"Error mounting Gradio app: {e}")
    raise

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)