import gradio as gr
from transformers import pipeline
from simple_salesforce import Salesforce
import datetime
import os
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Initialize Hugging Face model
generator = pipeline("text-generation", model="distilgpt2")

# Initialize Salesforce connection using environment variables
sf = Salesforce(
    username=os.getenv("SF_USERNAME"),
    password=os.getenv("SF_PASSWORD"),
    security_token=os.getenv("SF_SECURITY_TOKEN")
)

def generate_ai_data(supervisor_id, project_id, supervisor_data, project_data):
    """
    Generate AI coaching data and reports based on supervisor and project data.
    
    Args:
        supervisor_id (str): ID of the supervisor from Supervisor_Profile__c
        project_id (str): ID of the project from Project_Details__c
        supervisor_data (dict): Contains Role__c, Location__c
        project_data (dict): Contains Name, Start_Date__c, End_Date__c, Milestones__c, Project_Schedule__c
    
    Returns:
        dict: Status and generated data
    """
    try:
        # Construct prompt for AI generation
        prompt = (
            f"Generate daily checklist, tips, risk alerts, upcoming milestones, and performance trends for a "
            f"{supervisor_data['Role__c']} at {supervisor_data['Location__c']} working on project "
            f"{project_data['Name']} with milestones {project_data['Milestones__c']} and schedule "
            f"{project_data['Project_Schedule__c']}."
        )

        # Generate AI output
        ai_response = generator(prompt, max_length=500, num_return_sequences=1)[0]['generated_text']

        # Parse AI response (simplified parsing for this example)
        # In a real scenario, you'd use more sophisticated NLP to extract structured data
        daily_checklist = (
            "1. Conduct safety inspection of site (Safety, Pending)\n"
            "2. Ensure team wears protective gear (Safety, Pending)\n"
            "3. Schedule team briefing (General, Pending)"
        )
        suggested_tips = (
            "1. Prioritize safety checks due to upcoming weather risks.\n"
            "2. Focus on delayed tasks.\n"
            "3. Schedule a team review."
        )
        risk_alerts = "Risk of delay: Rain expected on May 22, 2025."
        upcoming_milestones = project_data['Milestones__c'].split(';')[0]  # Take the first milestone
        performance_trends = "Task completion rate: 75% this week (initial estimate)."

        # Save AI data to AI_Coaching_Data__c
        ai_data = {
            'Supervisor_ID__c': supervisor_id,
            'Project_ID__c': project_id,
            'Daily_Checklist__c': daily_checklist,
            'Suggested_Tips__c': suggested_tips,
            'Risk_Alerts__c': risk_alerts,
            'Upcoming_Milestones__c': upcoming_milestones,
            'Performance_Trends__c': performance_trends,
            'Generated_Date__c': datetime.datetime.now().strftime('%Y-%m-%d')
        }
        sf.AI_Coaching_Data__c.create(ai_data)

        # Generate a report for Report_Download__c
        report_data = {
            'Supervisor_ID__c': supervisor_id,
            'Project_ID__c': project_id,
            'Report_Type__c': 'Performance',
            'Report_Data__c': f"Performance Report: Task completion rate: 75% this week (initial estimate). Engagement score: 80%.",
            'Download_Link__c': 'https://salesforce-site.com/reports/RPT-0001.pdf',  # Update with actual Salesforce Site URL
            'Generated_Date__c': datetime.datetime.now().strftime('%Y-%m-%d')
        }
        sf.Report_Download__c.create(report_data)

        return {
            "status": "success",
            "message": "AI data and report generated successfully",
            "ai_data": ai_data,
            "report_data": report_data
        }

    except Exception as e:
        return {
            "status": "error",
            "message": f"Error generating AI data: {str(e)}"
        }

# Create Gradio interface
iface = gr.Interface(
    fn=generate_ai_data,
    inputs=[
        gr.Textbox(label="Supervisor ID"),
        gr.Textbox(label="Project ID"),
        gr.JSON(label="Supervisor Data"),
        gr.JSON(label="Project Data")
    ],
    outputs=gr.JSON(label="Result"),
    title="AI Coach Data Generator",
    description="Generate AI coaching data and reports based on supervisor and project details."
)

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)