from flask import Flask, request, jsonify
from transformers import pipeline
from simple_salesforce import Salesforce
import datetime
import os
from dotenv import load_dotenv

app = Flask(__name__)

# Load environment variables from .env file
load_dotenv()

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

# Initialize Salesforce connection using environment variables
try:
    sf = Salesforce(
        username=os.getenv("SF_USERNAME"),
        password=os.getenv("SF_PASSWORD"),
        security_token=os.getenv("SF_SECURITY_TOKEN")
    )
except Exception as e:
    print(f"Error connecting to Salesforce: {str(e)}")
    sf = None

@app.route('/generate-ai-data', methods=['POST'])
def generate_ai_data():
    """
    Generate AI coaching data and reports based on supervisor and project data.
    This endpoint is called by Salesforce when a new project is created or during daily refresh.
    """
    if sf is None:
        return jsonify({
            "status": "error",
            "message": "Salesforce connection failed. Check credentials in .env file."
        }), 500

    try:
        data = request.get_json()
        supervisor_id = data['supervisor_id']
        project_id = data['project_id']
        supervisor_data = data['supervisor_data']
        project_data = data['project_data']

        # Fetch existing checklists to determine progress
        checklists = sf.query(
            f"SELECT Task_Name__c, Status__c FROM Daily_Checklist__c WHERE Project_ID__c = '{project_id}' AND Date__c = TODAY"
        )['records']
        completed_tasks = [c['Task_Name__c'] for c in checklists if c['Status__c'] == 'Completed']
        pending_tasks = [c['Task_Name__c'] for c in checklists if c['Status__c'] == 'Pending']
        completion_rate = len(completed_tasks) / (len(completed_tasks) + len(pending_tasks)) * 100 if (len(completed_tasks) + len(pending_tasks)) > 0 else 0

        # 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']}. "
            f"Current completion rate: {completion_rate:.1f}%. Completed tasks: {', '.join(completed_tasks) if completed_tasks else 'None'}. "
            f"Pending tasks: {', '.join(pending_tasks) if pending_tasks else 'None'}."
        )

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

        # Parse AI response (more dynamic parsing based on project progress)
        today = datetime.datetime.now().strftime('%Y-%m-%d')
        daily_checklist = (
            f"1. Review pending tasks from yesterday (General, Pending)\n"
            f"2. Conduct daily safety inspection for {project_data['Name']} (Safety, Pending)\n"
            f"3. Schedule progress meeting (General, Pending)"
        ) if not pending_tasks else (
            f"1. Complete pending task: {pending_tasks[0]} (General, Pending)\n"
            f"2. Conduct daily safety inspection for {project_data['Name']} (Safety, Pending)\n"
            f"3. Schedule progress meeting (General, Pending)"
        )
        suggested_tips = (
            f"1. Focus on completing pending tasks to improve completion rate.\n"
            f"2. Monitor weather conditions in {supervisor_data['Location__c']}.\n"
            f"3. Prepare for upcoming milestone."
        )
        risk_alerts = f"Risk of delay: Weather risks in {supervisor_data['Location__c']} on {today}."
        upcoming_milestones = project_data['Milestones__c'].split(';')[0]  # Take the first milestone
        performance_trends = f"Task completion rate: {completion_rate:.1f}% (updated {today})."

        # 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': today
        }
        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': 'Daily Progress',
            'Report_Data__c': f"Daily Progress Report ({today}): Completion rate: {completion_rate:.1f}%. "
                             f"Pending tasks: {len(pending_tasks)}. Completed tasks: {len(completed_tasks)}.",
            'Download_Link__c': 'https://your-salesforce-site.com/reports/RPT-' + project_id + '.pdf',  # Update with actual Salesforce Site URL
            'Generated_Date__c': today
        }
        sf.Report_Download__c.create(report_data)

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

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

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