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

# Load environment variables from .env file (Salesforce creds, etc)
load_dotenv()

app = Flask(__name__)

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

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

@app.route('/generate_ai_data', methods=['POST'])
def generate_ai_data():
    try:
        data = request.json
        supervisor_id = data.get('supervisor_id')
        project_id = data.get('project_id')

        if not supervisor_id or not project_id:
            return jsonify({"status": "error", "message": "Missing supervisor_id or project_id"}), 400

        # Fetch supervisor data from Salesforce
        supervisor_data = sf.Supervisor_Profile__c.get(supervisor_id)
        # Fetch project data from Salesforce
        project_data = sf.Project_Details__c.get(project_id)

        # Construct prompt for AI
        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 response
        ai_response = generator(prompt, max_length=500, num_return_sequences=1)[0]['generated_text']

        # For demo: Simple static parsing; customize as needed
        daily_checklist = "1. Conduct safety inspection of site (Safety, Pending)"
        suggested_tips = "1. Prioritize safety checks due to upcoming weather risks."
        risk_alerts = "Risk of delay: Rain expected on May 22, 2025."
        upcoming_milestones = project_data['Milestones__c'].split(';')[0]
        performance_trends = "Task completion rate: 75% this week (initial estimate)."

        # Save AI data to Salesforce object 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)

        # Optionally generate a performance report and save to 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. Engagement score: 80%.",
            'Download_Link__c': 'https://salesforce-site.com/reports/RPT-0001.pdf',
            'Generated_Date__c': datetime.datetime.now().strftime('%Y-%m-%d')
        }
        sf.Report_Download__c.create(report_data)

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

    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500

if __name__ == "__main__":
    # Run Flask app on all interfaces, port 7860 (or your choice)
    app.run(host="0.0.0.0", port=7860)