app.py

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from simple_salesforce import Salesforce
import os
import base64
import datetime
from dotenv import load_dotenv
from fpdf import FPDF
import shutil
import html

# 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}")

# Load model and tokenizer
model_name = "distilgpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
model.config.pad_token_id = tokenizer.pad_token_id

# Prompt template
PROMPT_TEMPLATE = """You are an AI assistant for construction supervisors. Given the role, project, milestones, and a reflection log, generate:

1. A Daily Checklist with clear and concise tasks based on the role and milestones. 
   Split the checklist into day-by-day tasks for a specified time period (e.g., one week).
2. Focus Suggestions based on concerns or keywords in the reflection log. Provide at least 2 suggestions.

Inputs:
Role: {role}
Project ID: {project_id}
Milestones: {milestones}
Reflection Log: {reflection}

Output Format:
Checklist (Day-by-Day):
- Day 1: 
   - Task 1
   - Task 2
- Day 2: 
   - Task 1
   - Task 2
...
Suggestions:
- 
"""

# Function to clean the text for PDF
def clean_text_for_pdf(text):
    return html.unescape(text).encode('latin-1', 'replace').decode('latin-1')

# 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)

    # Copy the file to a temporary directory for Gradio to access
    temp_pdf_path = "/tmp/" + os.path.basename(file_path)  # Use /tmp/ or current directory for Gradio
    shutil.copy(file_path, temp_pdf_path)

    return temp_pdf_path, filename

# Upload to Salesforce and update record with the generated URL
def upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, 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')
        )
        
        # Read and encode the file as base64
        with open(pdf_path, "rb") as f:
            encoded = base64.b64encode(f.read()).decode()

        # Create ContentVersion record to upload the PDF to Salesforce
        content = sf.ContentVersion.create({
            'Title': pdf_name,
            'PathOnClient': pdf_name,
            'VersionData': encoded
        })

        # Get the ContentDocumentId for the uploaded PDF
        content_id = content['id']
        
        # Generate the download URL for the uploaded PDF
        download_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_id}"

        # Query Salesforce to find the specific Supervisor_AI_Coaching__c record
        query = sf.query(f"""
            SELECT Id FROM Supervisor_AI_Coaching__c 
            WHERE Project_ID__c = '{project_id}' 
            AND Name = '{supervisor_name}'
            LIMIT 1
        """)

        if query['totalSize'] > 0:
            # Get the ID of the Supervisor_AI_Coaching__c record
            coaching_id = query['records'][0]['Id']
            
            # Update the Supervisor_AI_Coaching__c record with the download URL
            sf.Supervisor_AI_Coaching__c.update(coaching_id, {
                'Download_Link__c': download_url  # Update the Download_Link__c field with the URL
            })

        return download_url

    except Exception as e:
        print(f"⚠️ Upload error: {e}")
        return ""

# Salesforce helpers
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:
        print(f"⚠️ Error fetching roles: {e}")
        return []

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:
        print(f"⚠️ Error fetching names: {e}")
        return []

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:
        print(f"⚠️ Error fetching project: {e}")
        return ""

# Generate Salesforce dashboard URL
def generate_salesforce_dashboard_url(supervisor_name, project_id):
    return f"https://aicoachforsitesupervisors-dev-ed--c.develop.vf.force.com/apex/DashboardPage?supervisorName={supervisor_name}&projectId={project_id}"

def open_dashboard(role, supervisor_name, project_id):
    url = generate_salesforce_dashboard_url(supervisor_name, project_id)
    return f'<a href="{url}" target="_blank">Open Salesforce Dashboard</a>'

# Generate AI output
def generate_outputs(role, supervisor_name, project_id, milestones, reflection):
    if not all([role, supervisor_name, project_id, milestones, reflection]):
        return "❗ Please fill all fields.", "", None, ""

    prompt = PROMPT_TEMPLATE.format(role=role, project_id=project_id, milestones=milestones, reflection=reflection)
    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
    try:
        with torch.no_grad():
            outputs = model.generate(
                inputs['input_ids'],
                max_new_tokens=150,
                no_repeat_ngram_size=2,
                do_sample=True,
                top_p=0.9,
                temperature=0.7,
                pad_token_id=tokenizer.pad_token_id
            )
        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    except Exception as e:
        print(f"⚠️ Generation error: {e}")
        return "", "", None, ""

    def extract_between(text, start, end):
        s = text.find(start)
        e = text.find(end, s) if end else len(text)
        return text[s + len(start):e].strip() if s != -1 else ""

    checklist = extract_between(result, "Checklist:\n", "Suggestions:")
    suggestions = extract_between(result, "Suggestions:\n", None)

    if not checklist.strip():
        checklist = "- Perform daily safety inspection"
    if not suggestions.strip():
        suggestions = "- Monitor team coordination\n- Review safety protocols with the team"

    pdf_path, pdf_name = save_report_as_pdf(role, supervisor_name, project_id, checklist, suggestions)
    pdf_url = upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, pdf_path, pdf_name)
    if pdf_url:
        suggestions += f"\n\n🔗 [Download PDF Report]({pdf_url})"

    return checklist, suggestions, pdf_path, pdf_path

# Gradio Interface
def create_interface():
    roles = get_roles_from_salesforce()
    with gr.Blocks(theme="soft") as demo:
        gr.Markdown("## 🧠 AI-Powered Supervisor Assistant")

        with gr.Row():
            role = gr.Dropdown(choices=roles, label="Role")
            supervisor_name = gr.Dropdown(choices=[], label="Supervisor Name")
            project_id = gr.Textbox(label="Project ID", interactive=False)

        milestones = gr.Textbox(label="Milestones (comma-separated KPIs)")
        reflection = gr.Textbox(label="Reflection Log", lines=4)

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

        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()
        dashboard_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, pdf_path, _ = generate_outputs(role, supervisor_name, project_id, milestones, reflection)
            return checklist, suggestions, pdf_path, pdf_path

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

        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)

        dashboard_btn.click(fn=open_dashboard,
                            inputs=[role, supervisor_name, project_id],
                            outputs=dashboard_link)

    return demo

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