app.py

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from simple_salesforce import Salesforce
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Check if required environment variables are set
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}")

# Get configurable values for KPI_Flag__c and Engagement_Score__c
KPI_FLAG_DEFAULT = os.getenv('KPI_FLAG', 'True') == 'True'  # Default to True if not set
ENGAGEMENT_SCORE_DEFAULT = float(os.getenv('ENGAGEMENT_SCORE', '85.0'))  # Default to 85.0

# Initialize model and tokenizer
model_name = "distilgpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Avoid warnings by setting pad token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token if tokenizer.eos_token else "[PAD]"
    tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
model.config.pad_token_id = tokenizer.pad_token_id

# Prompt template for generating structured output
PROMPT_TEMPLATE = """You are an AI coach for construction supervisors. Based on the following inputs, generate a daily checklist, focus suggestions, and a motivational quote.
Inputs:
Role: {role}
Project: {project_id}
Milestones: {milestones}
Reflection: {reflection}
Format your response clearly like this:
Checklist:
- {milestones_list}
Suggestions:
- {suggestions_list}
Quote:
- Your motivational quote here
"""

# Function to get all 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')
        )

        # Query distinct Role__c values
        result = sf.query("SELECT Role__c FROM Supervisor__c WHERE Role__c != NULL")

        # Extract roles and remove duplicates
        roles = list(set(record['Role__c'] for record in result.get('records', [])))

        print(f"✅ Fetched {len(roles)} unique roles from Salesforce")
        return roles

    except Exception as e:
        print(f"⚠️ Error fetching roles from Salesforce: {e}")
        print("Using fallback roles...")
        return ["Site Manager", "Safety Officer", "Project Lead"]  # Match actual active roles


# Function to get supervisor's Name (Auto Number) by 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')
        )

        # Escape single quotes in the role to prevent SOQL injection
        role = role.replace("'", "\\'")

        # Query all supervisors for the selected role
        result = sf.query(f"SELECT Name FROM Supervisor__c WHERE Role__c = '{role}'")
        if result['totalSize'] == 0:
            print("❌ No matching supervisors found.")
            return []

        # Extract all supervisor names
        supervisor_names = [record['Name'] for record in result['records']]
        print(f"✅ Found supervisors: {supervisor_names} for role: {role}")
        return supervisor_names

    except Exception as e:
        print(f"⚠️ Error fetching supervisor names: {e}")
        return []


# Function to get project IDs and names assigned to selected supervisor
def get_projects_for_supervisor(supervisor_name):
    try:
        # Use the selected supervisor name to fetch the associated project
        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')
        )

        # Escape single quotes in the supervisor_name
        supervisor_name = supervisor_name.replace("'", "\\'")

        # Step 1: Get the Salesforce record ID of the supervisor based on the Name
        supervisor_result = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
        if supervisor_result['totalSize'] == 0:
            print("❌ No supervisor found with the given name.")
            return ""

        supervisor_id = supervisor_result['records'][0]['Id']

        # Step 2: Query Project__c records where Supervisor_ID__c matches the supervisor's record ID
        project_result = sf.query(f"SELECT Name FROM Project__c WHERE Supervisor_ID__c = '{supervisor_id}' LIMIT 1")

        if project_result['totalSize'] == 0:
            print("❌ No project found for supervisor.")
            return ""

        project_name = project_result['records'][0]['Name']
        print(f"✅ Found project: {project_name} for supervisor: {supervisor_name}")
        return project_name

    except Exception as e:
        print(f"⚠️ Error fetching project for supervisor: {e}")
        return ""


# Function to generate AI-based coaching output
def generate_outputs(role, supervisor_name, project_id, milestones, reflection):
    if not all([role, supervisor_name, project_id, milestones, reflection]):
        return "Error: All fields are required.", "", ""

    # Format the prompt
    milestones_list = "\n- ".join([m.strip() for m in milestones.split(",")])

    suggestions_list = ""
    if "delays" in reflection.lower():
        suggestions_list = "- Consider adjusting timelines to accommodate delays.\n- Communicate delays to all relevant stakeholders."
    elif "weather" in reflection.lower():
        suggestions_list = "- Ensure team has rain gear.\n- Monitor weather updates for possible further delays."
    elif "equipment" in reflection.lower():
        suggestions_list = "- Inspect all equipment to ensure no malfunctions.\n- Schedule maintenance if necessary."

    # Fill in the prompt template
    prompt = PROMPT_TEMPLATE.format(
        role=role,
        project_id=project_id,
        milestones=milestones,
        reflection=reflection,
        milestones_list=milestones_list,
        suggestions_list=suggestions_list
    )

    # Tokenize input
    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True, padding=True)

    # Generate response
    try:
        with torch.no_grad():
            outputs = model.generate(
                inputs['input_ids'],
                max_length=1024,  # Increased to allow for longer outputs
                num_return_sequences=1,
                no_repeat_ngram_size=2,
                do_sample=True,
                top_p=0.9,
                temperature=0.8,
                pad_token_id=tokenizer.pad_token_id
            )

        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    except Exception as e:
        print(f"⚠️ Error during model generation: {e}")
        return "Error: Failed to generate outputs.", "", ""

    # Parse sections
    def extract_section(text, start_marker, end_marker):
        start = text.find(start_marker)
        if start == -1:
            return "Not found"
        start += len(start_marker)
        end = text.find(end_marker, start) if end_marker else len(text)
        return text[start:end].strip()

    checklist = extract_section(generated_text, "Checklist:\n", "Suggestions:")
    suggestions = extract_section(generated_text, "Suggestions:\n", "Quote:")
    quote = extract_section(generated_text, "Quote:\n", None)

    # Save to Salesforce
    save_to_salesforce(role, project_id, milestones, reflection, checklist, suggestions, quote, supervisor_name)

    return checklist, suggestions, quote


# Function to check if a field exists in a Salesforce object
def field_exists(sf, object_name, field_name):
    try:
        # Describe the object to get its fields
        obj_desc = getattr(sf, object_name).describe()
        fields = [field['name'] for field in obj_desc['fields']]
        return field_name in fields
    except Exception as e:
        print(f"⚠️ Error checking if field {field_name} exists in {object_name}: {e}")
        return False


# Function to create a record in Salesforce
def save_to_salesforce(role, project_id, milestones, reflection, checklist, suggestions, quote, 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')
        )

        # Escape single quotes in supervisor_name and project_id
        supervisor_name = supervisor_name.replace("'", "\\'")
        project_id = project_id.replace("'", "\\'")

        # Step 1: Get the Salesforce record ID for the supervisor
        supervisor_result = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
        if supervisor_result['totalSize'] == 0:
            print(f"❌ No supervisor found with Name: {supervisor_name}")
            return

        supervisor_id = supervisor_result['records'][0]['Id']

        # Step 2: Get the Salesforce record ID for the project
        project_result = sf.query(f"SELECT Id FROM Project__c WHERE Name = '{project_id}' LIMIT 1")
        if project_result['totalSize'] == 0:
            print(f"❌ No project found with Name: {project_id}")
            return

        project_record_id = project_result['records'][0]['Id']

        # Truncate text fields to avoid exceeding Salesforce field length limits (assuming 255 characters for simplicity)
        MAX_TEXT_LENGTH = 255
        checklist = checklist[:MAX_TEXT_LENGTH] if checklist else ""
        suggestions = suggestions[:MAX_TEXT_LENGTH] if suggestions else ""
        reflection = reflection[:MAX_TEXT_LENGTH] if reflection else ""

        # Prepare data for Salesforce with explicit mapping
        data = {
            'Supervisor_ID__c': supervisor_id,       # Lookup field expects the record ID of Supervisor__c
            'Project_ID__c': project_record_id,      # Lookup field expects the record ID of Project__c
            'Daily_Checklist__c': checklist,         # Maps to the generated Daily Checklist
            'Suggested_Tips__c': suggestions,        # Maps to the generated Focus Suggestions
            'Reflection_Log__c': reflection,         # Maps to the Reflection Log input
            'KPI_Flag__c': KPI_FLAG_DEFAULT,         # Configurable via .env
            'Engagement_Score__c': ENGAGEMENT_SCORE_DEFAULT  # Configurable via .env
        }

        # Check if Milestones_KPIs__c field exists before mapping
        if field_exists(sf, 'Supervisor_AI_Coaching__c', 'Milestones_KPIs__c'):
            # Truncate milestones as well if the field exists
            milestones = milestones[:MAX_TEXT_LENGTH] if milestones else ""
            data['Milestones_KPIs__c'] = milestones
        else:
            print("⚠️ Milestones_KPIs__c field does not exist in Supervisor_AI_Coaching__c. Skipping mapping.")

        # Create record
        response = sf.Supervisor_AI_Coaching__c.create(data)
        print("✅ Record created successfully in Salesforce.")
        print("Record ID:", response['id'])

    except Exception as e:
        print(f"❌ Error saving to Salesforce: {e}")
        print("Data being sent:", data)
        if hasattr(e, 'content'):
            print("Salesforce API response:", e.content)


# Gradio Interface
def create_interface():
    # Fetch roles from Salesforce
    roles = get_roles_from_salesforce()
    print(f"Fetched Roles: {roles}")

    with gr.Blocks(theme="soft") as demo:
        gr.Markdown("# 🏗️ Construction Supervisor AI Coach")
        gr.Markdown("Enter details to generate a daily checklist, focus suggestions, and a motivational quote.")

        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():
            submit = gr.Button("Generate", variant="primary")
            clear = gr.Button("Clear")
            refresh_btn = gr.Button("🔄 Refresh Roles")

        checklist_output = gr.Textbox(label="✅ Daily Checklist")
        suggestions_output = gr.Textbox(label="💡 Focus Suggestions")
        quote_output = gr.Textbox(label="✨ Motivational Quote")

        # Event: When role changes, update supervisor name dropdown
        role.change(
            fn=lambda r: gr.update(choices=get_supervisor_name_by_role(r)),
            inputs=[role],
            outputs=[supervisor_name]
        )

        # Event: When supervisor name changes, update project ID
        supervisor_name.change(
            fn=get_projects_for_supervisor,
            inputs=[supervisor_name],
            outputs=[project_id]
        )

        submit.click(
            fn=generate_outputs,
            inputs=[role, supervisor_name, project_id, milestones, reflection],
            outputs=[checklist_output, suggestions_output, quote_output]
        )

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

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

    return demo


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