import gradio as gr
import os
import base64
import datetime
from dotenv import load_dotenv
from fpdf import FPDF
import shutil
import html
import io
import matplotlib.pyplot as plt
import numpy as np
from simple_salesforce import Salesforce

# Load environment variables
load_dotenv()

# Required env vars check
required_env_vars = ['SF_USERNAME', 'SF_PASSWORD', 'SF_SECURITY_TOKEN']
missing_vars = [var for var in required_env_vars if not os.getenv(var)]
if missing_vars:
    raise EnvironmentError(f"Missing required environment variables: {missing_vars}")

# Defaults
KPI_FLAG_DEFAULT = os.getenv('KPI_FLAG', 'True') == 'True'
ENGAGEMENT_SCORE_DEFAULT = float(os.getenv('ENGAGEMENT_SCORE', '85.0'))

# Function to generate progress chart
def show_dashboard_chart(start_date, end_date, tasks_completed):
    completed_tasks = list(tasks_completed.values())
    labels = list(tasks_completed.keys())
    remaining_tasks = [5 - task for task in completed_tasks]  # Assuming 5 tasks per date

    fig, ax = plt.subplots(figsize=(10, 6))
    ax.bar(labels, completed_tasks, color='green', label="Completed")
    ax.bar(labels, remaining_tasks, bottom=completed_tasks, color='gray', label="Remaining")

    ax.set_xlabel("Dates")
    ax.set_ylabel("Task Progress")
    ax.set_title(f"Task Completion Progress from {start_date} to {end_date}")
    ax.legend()
    plt.xticks(rotation=45)
    chart_image = io.BytesIO()
    plt.savefig(chart_image, format='png')
    chart_image.seek(0)
    return chart_image

# Function to get the data from Salesforce
def get_dashboard_data_from_salesforce(supervisor_name, project_id):
    try:
        sf = Salesforce(
            username=os.getenv('SF_USERNAME'),
            password=os.getenv('SF_PASSWORD'),
            security_token=os.getenv('SF_SECURITY_TOKEN'),
            domain=os.getenv('SF_DOMAIN', 'login')
        )

        query = sf.query(f"SELECT Start_Date__c, End_Date__c FROM Project__c WHERE Name = '{project_id}' LIMIT 1")
        if query['totalSize'] == 0:
            return "", "", None, "Project not found"

        start_date_str = query['records'][0]['Start_Date__c']
        end_date_str = query['records'][0]['End_Date__c']

        start_date = datetime.datetime.strptime(start_date_str, "%Y-%m-%d")
        end_date = datetime.datetime.strptime(end_date_str, "%Y-%m-%d")

        task_dates = [start_date + datetime.timedelta(days=i) for i in range((end_date - start_date).days + 1)]
        tasks_completed = {str(task_dates[i].date()): np.random.randint(1, 6) for i in range(len(task_dates))}

        chart_image = show_dashboard_chart(start_date, end_date, tasks_completed)
        return start_date, end_date, chart_image, f"Project {project_id} Task Progress"
    except Exception as e:
        return "", "", None, f"Error: {str(e)}"

# Clean text for PDF generation
def clean_text_for_pdf(text):
    return html.unescape(text).encode('latin-1', 'replace').decode('latin-1')

# Function to save report as PDF
def save_report_as_pdf(role, supervisor_name, project_id, checklist, suggestions):
    now = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    filename = f"report_{supervisor_name}_{project_id}_{now}.pdf"
    file_path = f"./reports/{filename}"
    os.makedirs("reports", exist_ok=True)

    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", 'B', 14)
    pdf.cell(200, 10, txt="Supervisor Daily Report", ln=True, align="C")
    pdf.set_font("Arial", size=12)
    pdf.cell(200, 10, txt=clean_text_for_pdf(f"Role: {role}"), ln=True)
    pdf.cell(200, 10, txt=clean_text_for_pdf(f"Supervisor: {supervisor_name}"), ln=True)
    pdf.cell(200, 10, txt=clean_text_for_pdf(f"Project ID: {project_id}"), ln=True)
    pdf.ln(5)
    pdf.set_font("Arial", 'B', 12)
    pdf.cell(200, 10, txt="Daily Checklist", ln=True)
    pdf.set_font("Arial", size=12)
    for line in checklist.split("\n"):
        pdf.multi_cell(0, 10, clean_text_for_pdf(line))
    pdf.ln(5)
    pdf.set_font("Arial", 'B', 12)
    pdf.cell(200, 10, txt="Focus Suggestions", ln=True)
    pdf.set_font("Arial", size=12)
    for line in suggestions.split("\n"):
        pdf.multi_cell(0, 10, clean_text_for_pdf(line))
    pdf.output(file_path)

    temp_pdf_path = "/tmp/" + os.path.basename(file_path)
    shutil.copy(file_path, temp_pdf_path)
    return temp_pdf_path, filename

# Function to get roles from Salesforce
def get_roles_from_salesforce():
    try:
        sf = Salesforce(
            username=os.getenv('SF_USERNAME'),
            password=os.getenv('SF_PASSWORD'),
            security_token=os.getenv('SF_SECURITY_TOKEN'),
            domain=os.getenv('SF_DOMAIN', 'login')
        )
        result = sf.query("SELECT Role__c FROM Supervisor__c WHERE Role__c != NULL")
        return list(set(record['Role__c'] for record in result['records']))
    except Exception as e:
        return []

# Function to get supervisor names based on role
def get_supervisor_name_by_role(role):
    try:
        sf = Salesforce(
            username=os.getenv('SF_USERNAME'),
            password=os.getenv('SF_PASSWORD'),
            security_token=os.getenv('SF_SECURITY_TOKEN'),
            domain=os.getenv('SF_DOMAIN', 'login')
        )
        result = sf.query(f"SELECT Name FROM Supervisor__c WHERE Role__c = '{role}'")
        return [record['Name'] for record in result['records']]
    except Exception as e:
        return []

# Function to get the project name for a supervisor
def get_projects_for_supervisor(supervisor_name):
    try:
        sf = Salesforce(
            username=os.getenv('SF_USERNAME'),
            password=os.getenv('SF_PASSWORD'),
            security_token=os.getenv('SF_SECURITY_TOKEN'),
            domain=os.getenv('SF_DOMAIN', 'login')
        )
        result = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
        if result['totalSize'] == 0:
            return ""
        supervisor_id = result['records'][0]['Id']
        project_result = sf.query(f"SELECT Name FROM Project__c WHERE Supervisor_ID__c = '{supervisor_id}' LIMIT 1")
        return project_result['records'][0]['Name'] if project_result['totalSize'] > 0 else ""
    except Exception as e:
        return ""

# Function to generate daily checklist and focus suggestions
def generate_checklist_and_suggestions(role, project_id, milestones, reflection):
    if not all([role, project_id, milestones, reflection]):
        return "Please fill all fields.", ""

    # Generate Checklist
    checklist_items = []
    if role.lower() == "site supervisor":
        checklist_items.append("- Conduct daily site safety briefing")
        checklist_items.append("- Monitor progress against project timeline")

    milestones_lower = milestones.lower()
    if "excavation" in milestones_lower:
        checklist_items.append("- Verify excavation progress and safety measures")
    if "foundation work" in milestones_lower:
        checklist_items.append("- Inspect foundation work for quality")
    if "pour the foundation concrete" in milestones_lower:
        checklist_items.append("- Ensure proper concrete pouring and curing")
    if "set up reinforcement bars" in milestones_lower:
        checklist_items.append("- Check reinforcement bar installation")
    if "test soil compaction" in milestones_lower:
        checklist_items.append("- Conduct soil compaction tests")

    if not checklist_items:
        checklist_items.append("- Perform general site inspection")

    checklist = "\n".join(checklist_items)

    # Generate Suggestions
    suggestions_items = []
    reflection_lower = reflection.lower()
    if "weather conditions" in reflection_lower or "rainfall" in reflection_lower:
        suggestions_items.append("- Adjust schedule for weather delays")
        suggestions_items.append("- Implement weather protection measures")
    if "delays" in reflection_lower:
        suggestions_items.append("- Identify causes of delays and mitigate")
    if "leveling" in reflection_lower or "base" in reflection_lower:
        suggestions_items.append("- Verify leveling accuracy and resolve issues")

    if not suggestions_items:
        suggestions_items.append("- Monitor team performance")
        suggestions_items.append("- Review safety protocols")

    suggestions = "\n".join(suggestions_items)

    return checklist, suggestions

# Function to upload the report and create the Supervisor AI Coaching record in Salesforce
def upload_report_and_create_supervisor_ai_coaching(supervisor_name, project_id, checklist, suggestions, pdf_path, pdf_name):
    try:
        sf = Salesforce(
            username=os.getenv('SF_USERNAME'),
            password=os.getenv('SF_PASSWORD'),
            security_token=os.getenv('SF_SECURITY_TOKEN'),
            domain=os.getenv('SF_DOMAIN', 'login')
        )

        with open(pdf_path, "rb") as f:
            encoded = base64.b64encode(f.read()).decode()

        content = sf.ContentVersion.create({
            'Title': pdf_name,
            'PathOnClient': pdf_name,
            'VersionData': encoded
        })

        content_id = content['id']
        download_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_id}"

        query = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
        supervisor_id = query['records'][0]['Id'] if query['totalSize'] > 0 else None
        if not supervisor_id:
            return "Supervisor not found."

        project_query = sf.query(f"SELECT Id FROM Project__c WHERE Name = '{project_id}' LIMIT 1")
        project_id_sf = project_query['records'][0]['Id'] if project_query['totalSize'] > 0 else None
        if not project_id_sf:
            return "Project not found."

        sf.Supervisor_AI_Coaching__c.create({
            'Project_ID__c': project_id_sf,
            'Supervisor_ID__c': supervisor_id,
            'Daily_Checklist__c': checklist,
            'Suggested_Tips__c': suggestions,
            'Download_Link__c': download_url
        })

        return "Supervisor AI Coaching record created and report uploaded successfully."
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio interface
def create_interface():
    roles = get_roles_from_salesforce()
    with gr.Blocks(theme="soft", css=".footer { display: none; }") as demo:
        gr.Markdown("## 🧠 AI-Powered Supervisor Assistant")

        with gr.Row():
            role = gr.Dropdown(choices=roles, label="Role", value="Site Supervisor")
            supervisor_name = gr.Dropdown(choices=["SUP-056"], label="Supervisor Name", value="SUP-056")
            project_id = gr.Textbox(label="Project ID", value="PROJ-078", interactive=False)

        milestones = gr.Textbox(label="Milestones (comma-separated KPIs)", value="Excavation and Foundation Work, Pour the foundation concrete, Set up reinforcement bars for the foundation, Test soil compaction")
        reflection = gr.Textbox(label="Reflection Log", lines=4, value="The excavation and foundation work was slower than expected due to unexpected weather conditions. We faced some delays because of heavy rainfall. The foundation team had some minor issues with leveling the base but they managed to resolve it in the end.")

        with gr.Row():
            generate = gr.Button("Generate")
            clear = gr.Button("Clear")
            refresh = gr.Button("🔄 Refresh Roles")

        checklist_output = gr.Textbox(label="✅ Daily Checklist")
        suggestions_output = gr.Textbox(label="💡 Focus Suggestions")
        download_button = gr.File(label="⬇ Download Report")
        pdf_link = gr.HTML()

        role.change(fn=lambda r: gr.update(choices=get_supervisor_name_by_role(r)), inputs=role, outputs=supervisor_name)
        supervisor_name.change(fn=get_projects_for_supervisor, inputs=supervisor_name, outputs=project_id)

        def handle_generate(role, supervisor_name, project_id, milestones, reflection):
            checklist, suggestions = generate_checklist_and_suggestions(role, project_id, milestones, reflection)
            if not checklist or not suggestions:
                return "Error generating checklist or suggestions.", "Error generating checklist or suggestions.", None, None, "Error: Generation failed."
            pdf_path, pdf_name = save_report_as_pdf(role, supervisor_name, project_id, checklist, suggestions)
            supervisor_ai_coaching_response = upload_report_and_create_supervisor_ai_coaching(supervisor_name, project_id, checklist, suggestions, pdf_path, pdf_name)
            pdf_link_content = f'<a href="file={pdf_path}" download="{pdf_name}">Download PDF</a>'
            return checklist, suggestions, pdf_path, pdf_link_content, supervisor_ai_coaching_response

        generate.click(fn=handle_generate,
                       inputs=[role, supervisor_name, project_id, milestones, reflection],
                       outputs=[checklist_output, suggestions_output, download_button, pdf_link, gr.HTML()])

        clear.click(fn=lambda: ("", "", "", "", ""),
                    inputs=None,
                    outputs=[role, supervisor_name, project_id, milestones, reflection])

        refresh.click(fn=lambda: gr.update(choices=get_roles_from_salesforce()), outputs=role)

        with gr.Tab("📊 Supervisor Dashboard"):
            dash_supervisor = gr.Textbox(label="Supervisor Name", placeholder="e.g., SUP-056")
            dash_project = gr.Textbox(label="Project ID", placeholder="e.g., PROJ-078")
            load_dash = gr.Button("📥 Load Dashboard")
            dash_output = gr.HTML()

            def show_dashboard_html(sup_name, proj_id):
                start_date, end_date, chart_image, chart_title = get_dashboard_data_from_salesforce(sup_name, proj_id)
                if chart_image:
                    chart_html = f"<img src='data:image/png;base64,{base64.b64encode(chart_image.read()).decode()}' />"
                    return f"<h3>{chart_title}</h3><p>From {start_date} to {end_date}</p>{chart_html}"
                else:
                    return f"Error: {chart_title}"

            load_dash.click(fn=show_dashboard_html, inputs=[dash_supervisor, dash_project], outputs=dash_output)

    return demo

if __name__ == "__main__":
    app = create_interface()
    app.launch()