app.py

from flask import Flask, request, jsonify, session, redirect, url_for, render_template, send_file
from simple_salesforce import Salesforce, SalesforceAuthenticationFailed, SalesforceError
import os
from datetime import datetime
from dotenv import load_dotenv
import logging
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
import io
import tempfile

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('app.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

app = Flask(__name__)
app.secret_key = os.getenv('FLASK_SECRET_KEY', 'your-secret-key')

# Salesforce mock data for guest users
MOCK_DATA = {
    "supervisor_name": "GUEST",
    "project_id": "PROJ_001",
    "daily_checklist": "Inspect safety equipment\nReview team assignments\nCheck project timeline",
    "suggested_tips": "Prioritize safety checks\nCommunicate clearly with the team\nMonitor resource usage",
    "reflection_log": "",
    "engagement_score": 85,
    "kpi_flag": False,
    "download_link": "",
    "last_login": str(datetime.now())
}

def get_salesforce_connection():
    """Establish a Salesforce connection with error handling."""
    try:
        sf = Salesforce(
            username=os.getenv('SALESFORCE_USERNAME'),
            password=os.getenv('SALESFORCE_PASSWORD'),
            security_token=os.getenv('SALESFORCE_SECURITY_TOKEN'),
            domain=os.getenv('SALESFORCE_DOMAIN', 'test')  # 'login' for production
        )
        logger.info("Successfully connected to Salesforce")
        return sf
    except SalesforceAuthenticationFailed as e:
        logger.error(f"Salesforce authentication failed: {str(e)}")
        raise Exception("Salesforce authentication failed. Check your credentials.")
    except Exception as e:
        logger.error(f"Error connecting to Salesforce: {str(e)}")
        raise Exception("Unable to connect to Salesforce. Please try again later.")

@app.route('/')
def index():
    if 'supervisor_name' not in session:
        logger.info("User not logged in, redirecting to login page")
        return redirect(url_for('login_page'))
    return render_template('index.html')

@app.route('/login', methods=['GET'])
def login_page():
    return render_template('login.html')

@app.route('/signup', methods=['GET'])
def signup_page():
    return render_template('signup.html')

@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')  # Now using Name field
    password = data.get('password')

    if username == 'GUEST':
        session['supervisor_name'] = 'GUEST'
        logger.info("Guest login successful")
        return jsonify({"status": "success", "message": "Logged in as guest"})

    try:
        sf = get_salesforce_connection()
        # Query Supervisor_AI_Coaching__c for Name and Password__c
        query = f"SELECT Id, Name, Password__c FROM Supervisor_AI_Coaching__c WHERE Name = '{username}' LIMIT 1"
        result = sf.query(query)
        if not result['records']:
            logger.warning(f"Invalid username: {username}")
            return jsonify({"status": "error", "message": "Invalid username"}), 401

        record = result['records'][0]
        stored_password = record['Password__c']
        if stored_password != password:
            logger.warning(f"Invalid password for username: {username}")
            return jsonify({"status": "error", "message": "Invalid password"}), 401

        session['supervisor_name'] = username
        logger.info(f"Login successful for {username}")
        return jsonify({"status": "success", "message": "Login successful"})
    except Exception as e:
        logger.error(f"Login error: {str(e)}")
        return jsonify({"status": "error", "message": str(e)}), 500

@app.route('/signup', methods=['POST'])
def signup():
    data = request.get_json()
    username = data.get('username')  # Name field
    password = data.get('password')  # Password__c field
    project_id = data.get('project_id', 'PROJ_001')  # Project_ID__c
    engagement_score = float(data.get('engagement_score', 85))  # Engagement_Score__c
    kpi_flag = data.get('kpi_flag', False)  # KPI_Flag__c

    if not username or not password:
        logger.warning("Signup failed: Username and password are required")
        return jsonify({"status": "error", "message": "Username and password are required"}), 400

    try:
        sf = get_salesforce_connection()
        # Check if username already exists
        query = f"SELECT Id FROM Supervisor_AI_Coaching__c WHERE Name = '{username}' LIMIT 1"
        result = sf.query(query)
        if result['records']:
            logger.warning(f"Signup failed: Username {username} already exists")
            return jsonify({"status": "error", "message": "Username already exists"}), 400

        # Create new Supervisor_AI_Coaching__c record
        new_record = {
            'Name': username,
            'Password__c': password,
            'Project_ID__c': project_id,
            'Engagement_Score__c': engagement_score,
            'KPI_Flag__c': kpi_flag,
            'Daily_Checklist__c': '',
            'Suggested_Tips__c': '',
            'Reflection_Log__c': '',
            'Download_Link__c': ''
        }
        sf.Supervisor_AI_Coaching__c.create(new_record)
        logger.info(f"Signup successful for {username}")

        # Automatically log in the user after signup
        session['supervisor_name'] = username
        return jsonify({"status": "success", "message": "Signup successful, you are now logged in"})
    except SalesforceError as e:
        logger.error(f"Salesforce API error during signup: {str(e)}")
        return jsonify({"status": "error", "message": "Salesforce API error. Please try again later."}), 500
    except Exception as e:
        logger.error(f"Signup error: {str(e)}")
        return jsonify({"status": "error", "message": str(e)}), 500

@app.route('/logout', methods=['POST'])
def logout():
    supervisor_name = session.get('supervisor_name', 'Unknown')
    session.pop('supervisor_name', None)
    logger.info(f"User {supervisor_name} logged out")
    return jsonify({"status": "success", "message": "Logged out successfully"})

@app.route('/get_supervisor_data')
def get_supervisor_data():
    supervisor_name = session.get('supervisor_name', 'GUEST')
    if supervisor_name == 'GUEST':
        logger.info("Returning mock data for guest user")
        return jsonify({"status": "success", "data": MOCK_DATA})

    try:
        sf = get_salesforce_connection()
        query = f"""
            SELECT Name, Project_ID__c, Daily_Checklist__c, Suggested_Tips__c,
                   Reflection_Log__c, Engagement_Score__c, KPI_Flag__c, Download_Link__c
            FROM Supervisor_AI_Coaching__c
            WHERE Name = '{supervisor_name}'
            LIMIT 1
        """
        result = sf.query(query)

        if result['records']:
            record = result['records'][0]
            data = {
                "supervisor_name": record['Name'],
                "project_id": record['Project_ID__c'],
                "daily_checklist": record['Daily_Checklist__c'] or "",
                "suggested_tips": record['Suggested_Tips__c'] or "",
                "reflection_log": record['Reflection_Log__c'] or "",
                "engagement_score": record['Engagement_Score__c'] or 0,
                "kpi_flag": record['KPI_Flag__c'],
                "download_link": record['Download_Link__c'] or "",
                "last_login": str(datetime.now())
            }
            logger.info(f"Fetched data for supervisor {supervisor_name}")
            return jsonify({"status": "success", "data": data})
        else:
            logger.warning(f"No data found for supervisor {supervisor_name}")
            return jsonify({"status": "error", "message": "No data found for this supervisor"})
    except SalesforceError as e:
        logger.error(f"Salesforce API error while fetching data: {str(e)}")
        return jsonify({"status": "error", "message": "Salesforce API error. Please try again later."}), 500
    except Exception as e:
        logger.error(f"Error fetching supervisor data: {str(e)}")
        return jsonify({"status": "error", "message": str(e)}), 500

@app.route('/submit_reflection', methods=['POST'])
def submit_reflection():
    supervisor_name = session.get('supervisor_name', 'GUEST')
    if supervisor_name == 'GUEST':
        MOCK_DATA['reflection_log'] = request.get_json().get('reflection')
        logger.info("Reflection submitted for guest user")
        return jsonify({"status": "success", "message": "Reflection submitted successfully (guest mode)"})

    data = request.get_json()
    reflection = data.get('reflection')

    try:
        sf = get_salesforce_connection()
        query = f"SELECT Id FROM Supervisor_AI_Coaching__c WHERE Name = '{supervisor_name}' LIMIT 1"
        result = sf.query(query)
        if not result['records']:
            logger.warning(f"No record found for supervisor {supervisor_name}")
            return jsonify({"status": "error", "message": "No record found for this supervisor"}), 404

        record_id = result['records'][0]['Id']
        sf.Supervisor_AI_Coaching__c.update(record_id, {'Reflection_Log__c': reflection})
        logger.info(f"Reflection updated for supervisor {supervisor_name}")
        return jsonify({"status": "success", "message": "Reflection submitted successfully"})
    except SalesforceError as e:
        logger.error(f"Salesforce API error while submitting reflection: {str(e)}")
        return jsonify({"status": "error", "message": "Salesforce API error. Please try again later."}), 500
    except Exception as e:
        logger.error(f"Error submitting reflection: {str(e)}")
        return jsonify({"status": "error", "message": str(e)}), 500

@app.route('/generate', methods=['POST'])
def generate():
    supervisor_name = session.get('supervisor_name', 'GUEST')
    data = request.get_json()

    # Generate checklist and tips (simplified logic)
    checklist = [
        "Inspect safety equipment",
        "Review team assignments",
        "Check project timeline"
    ]
    tips = [
        "Prioritize safety checks",
        "Communicate clearly with the team",
        "Monitor resource usage"
    ]

    if supervisor_name == 'GUEST':
        MOCK_DATA['daily_checklist'] = '\n'.join(checklist)
        MOCK_DATA['suggested_tips'] = '\n'.join(tips)
        logger.info("Generated coaching output for guest user")
        return jsonify({
            "status": "success",
            "output": {
                "checklist": checklist,
                "tips": tips
            }
        })

    try:
        sf = get_salesforce_connection()
        query = f"SELECT Id FROM Supervisor_AI_Coaching__c WHERE Name = '{supervisor_name}' LIMIT 1"
        result = sf.query(query)
        if not result['records']:
            logger.warning(f"No record found for supervisor {supervisor_name}")
            return jsonify({"status": "error", "message": "No record found for this supervisor"}), 404

        record_id = result['records'][0]['Id']
        sf.Supervisor_AI_Coaching__c.update(record_id, {
            'Daily_Checklist__c': '\n'.join(checklist),
            'Suggested_Tips__c': '\n'.join(tips)
        })
        logger.info(f"Generated and updated coaching output for supervisor {supervisor_name}")
        return jsonify({
            "status": "success",
            "output": {
                "checklist": checklist,
                "tips": tips
            }
        })
    except SalesforceError as e:
        logger.error(f"Salesforce API error while generating output: {str(e)}")
        return jsonify({"status": "error", "message": "Salesforce API error. Please try again later."}), 500
    except Exception as e:
        logger.error(f"Error generating coaching output: {str(e)}")
        return jsonify({"status": "error", "message": str(e)}), 500

def generate_pdf_summary(data):
    """Generate a PDF summary and return the file path."""
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf', dir='/tmp') as tmp_file:
            pdf_path = tmp_file.name
            c = canvas.Canvas(pdf_path, pagesize=letter)
            width, height = letter

            # Title
            c.setFont("Helvetica-Bold", 16)
            c.drawString(50, height - 50, "Supervisor AI Coaching Summary")

            # Supervisor Data
            c.setFont("Helvetica", 12)
            y_position = height - 100
            c.drawString(50, y_position, f"Supervisor Name: {data.get('supervisor_name', 'N/A')}")
            y_position -= 20
            c.drawString(50, y_position, f"Project ID: {data.get('project_id', 'N/A')}")
            y_position -= 20
            c.drawString(50, y_position, f"Last Login: {data.get('last_login', 'N/A')}")
            y_position -= 40

            # Daily Checklist
            c.setFont("Helvetica-Bold", 14)
            c.drawString(50, y_position, "Daily Checklist")
            c.setFont("Helvetica", 12)
            y_position -= 20
            checklist = data.get('daily_checklist', '').split('\n')
            for item in checklist:
                if item.strip():
                    c.drawString(70, y_position, f"- {item}")
                    y_position -= 20
                    if y_position < 50:
                        c.showPage()
                        y_position = height - 50

            # Suggested Tips
            y_position -= 20
            c.setFont("Helvetica-Bold", 14)
            c.drawString(50, y_position, "Suggested Tips")
            c.setFont("Helvetica", 12)
            y_position -= 20
            tips = data.get('suggested_tips', '').split('\n')
            for tip in tips:
                if tip.strip():
                    c.drawString(70, y_position, f"- {tip}")
                    y_position -= 20
                    if y_position < 50:
                        c.showPage()
                        y_position = height - 50

            # Reflection Log
            y_position -= 20
            c.setFont("Helvetica-Bold", 14)
            c.drawString(50, y_position, "Reflection Log")
            c.setFont("Helvetica", 12)
            y_position -= 20
            reflection = data.get('reflection_log', 'No reflection available.')
            lines = reflection.split('\n')
            for line in lines:
                c.drawString(70, y_position, line)
                y_position -= 20
                if y_position < 50:
                    c.showPage()
                    y_position = height - 50

            # KPIs
            y_position -= 20
            c.setFont("Helvetica-Bold", 14)
            c.drawString(50, y_position, "KPIs")
            c.setFont("Helvetica", 12)
            y_position -= 20
            c.drawString(70, y_position, f"Engagement Score: {data.get('engagement_score', 0)}%")
            y_position -= 20
            c.drawString(70, y_position, f"KPI Flag: {'Active' if data.get('kpi_flag', False) else 'Inactive'}")
            y_position -= 20

            c.save()
            logger.info(f"Generated PDF at {pdf_path}")
            return pdf_path
    except Exception as e:
        logger.error(f"Error generating PDF: {str(e)}")
        raise

@app.route('/download_pdf', methods=['GET'])
def download_pdf():
    supervisor_name = session.get('supervisor_name', 'GUEST')
    if supervisor_name == 'GUEST':
        logger.info("Download not available for guest user")
        return jsonify({"status": "success", "message": "Download not available in guest mode"})

    try:
        sf = get_salesforce_connection()
        query = f"""
            SELECT Id, Name, Project_ID__c, Daily_Checklist__c, Suggested_Tips__c,
                   Reflection_Log__c, Engagement_Score__c, KPI_Flag__c, Download_Link__c
            FROM Supervisor_AI_Coaching__c
            WHERE Name = '{supervisor_name}'
            LIMIT 1
        """
        result = sf.query(query)
        if not result['records']:
            logger.warning(f"No record found for supervisor {supervisor_name}")
            return jsonify({"status": "error", "message": "No record found for this supervisor"}), 404

        record = result['records'][0]
        record_id = record['Id']
        data = {
            "supervisor_name": record['Name'],
            "project_id": record['Project_ID__c'],
            "daily_checklist": record['Daily_Checklist__c'] or "",
            "suggested_tips": record['Suggested_Tips__c'] or "",
            "reflection_log": record['Reflection_Log__c'] or "",
            "engagement_score": record['Engagement_Score__c'] or 0,
            "kpi_flag": record['KPI_Flag__c'],
            "download_link": record['Download_Link__c'] or "",
            "last_login": str(datetime.now())
        }

        # Generate PDF
        pdf_path = generate_pdf_summary(data)
        
        # Placeholder for Download_Link__c (requires actual storage solution in production)
        download_link = "https://example.com/report.pdf"
        sf.Supervisor_AI_Coaching__c.update(record_id, {'Download_Link__c': download_link})
        logger.info(f"Updated Download_Link__c for supervisor {supervisor_name}")

        # Serve the PDF file
        return send_file(
            pdf_path,
            as_attachment=True,
            download_name=f"supervisor_report_{supervisor_name}.pdf",
            mimetype='application/pdf'
        )
    except SalesforceError as e:
        logger.error(f"Salesforce API error while downloading PDF: {str(e)}")
        return jsonify({"status": "error", "message": "Salesforce API error. Please try again later."}), 500
    except Exception as e:
        logger.error(f"Error downloading PDF: {str(e)}")
        return jsonify({"status": "error", "message": str(e)}), 500
    finally:
        if 'pdf_path' in locals():
            try:
                os.remove(pdf_path)
                logger.info(f"Cleaned up temporary PDF file: {pdf_path}")
            except Exception as e:
                logger.warning(f"Failed to clean up PDF file: {str(e)}")

if __name__ == '__main__':
    port = int(os.getenv('PORT', 5000))
    app.run(host='0.0.0.0', port=port, debug=True)