from fastapi import FastAPI, HTTPException, Header
from pydantic import BaseModel
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
import base64
import os
import logging
from datetime import datetime
from fastapi.responses import HTMLResponse
from simple_salesforce import Salesforce
from dotenv import load_dotenv
from datasets import load_dataset

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

app = FastAPI()

# Load environment variables
load_dotenv()

# Environment variables
SF_USERNAME = os.getenv("SF_USERNAME")
SF_PASSWORD = os.getenv("SF_PASSWORD")
SF_SECURITY_TOKEN = os.getenv("SF_SECURITY_TOKEN")
SF_DOMAIN = os.getenv("SF_DOMAIN", "login")
API_KEY = os.getenv("HUGGINGFACE_API_KEY")

# Validate environment variables
required_env_vars = ["SF_USERNAME", "SF_PASSWORD", "SF_SECURITY_TOKEN", "HUGGINGFACE_API_KEY"]
for var in required_env_vars:
    if not os.getenv(var):
        logger.error(f"Environment variable {var} is not set")
        raise ValueError(f"Environment variable {var} is not set")

# Salesforce connection
sf = None
try:
    sf = Salesforce(
        username=SF_USERNAME,
        password=SF_PASSWORD,
        security_token=SF_SECURITY_TOKEN,
        domain=SF_DOMAIN
    )
    logger.info("Successfully connected to Salesforce")
except Exception as e:
    logger.error(f"Failed to connect to Salesforce: {str(e)}")
    raise RuntimeError(f"Cannot connect to Salesforce: {str(e)}")

# VendorLog model
class VendorLog(BaseModel):
    vendorLogId: str
    vendorId: str
    vendorRecordId: str
    workDetails: str
    qualityReport: str
    incidentLog: str
    workCompletionDate: str
    actualCompletionDate: str
    vendorLogName: str
    delayDays: int
    project: str

# Store vendor logs
vendor_logs = []

def validate_salesforce_fields():
    """Validate required Salesforce fields, log warnings if optional fields are missing."""
    try:
        vendor_log_fields = [f['name'] for f in sf.Vendor_Log__c.describe()['fields']]
        required_vendor_fields = [
            'Vendor__c', 'Work_Completion_Percentage__c', 'Quality_Percentage__c',
            'Incident_Severity__c', 'Work_Completion_Date__c', 'Actual_Completion_Date__c',
            'Delay_Days__c', 'Project__c'
        ]
        for field in required_vendor_fields:
            if field not in vendor_log_fields:
                logger.error(f"Field {field} not found in Vendor_Log__c")
                raise ValueError(f"Field {field} not found in Vendor_Log__c")

        score_fields = [f['name'] for f in sf.Subcontractor_Performance_Score__c.describe()['fields']]
        required_score_fields = [
            'Vendor_Log__c', 'Vendor__c', 'Quality_Score__c', 'Timeliness_Score__c',
            'Safety_Score__c', 'Communication_Score__c', 'Alert_Flag__c'
        ]
        optional_score_fields = ['PDF_Link__c']
        for field in required_score_fields:
            if field not in score_fields:
                logger.error(f"Field {field} not found in Subcontractor_Performance_Score__c")
                raise ValueError(f"Field {field} not found in Subcontractor_Performance_Score__c")
        for field in optional_score_fields:
            if field not in score_fields:
                logger.warning(f"Optional field {field} not found in Subcontractor_Performance_Score__c, skipping related functionality")
        logger.info("Salesforce fields validated successfully")
        return 'PDF_Link__c' in score_fields
    except Exception as e:
        logger.error(f"Error validating Salesforce fields: {str(e)}")
        raise

# Check if PDF_Link__c exists
has_pdf_link = validate_salesforce_fields()

def fetch_huggingface_records(dataset_name: str = "imdb"):
    """Fetch records from a Hugging Face dataset."""
    try:
        os.environ["HUGGINGFACE_TOKEN"] = API_KEY
        dataset = load_dataset(dataset_name)
        logger.info(f"Successfully fetched dataset: {dataset_name}")
        records = [record for record in dataset['train']]  # Assuming 'train' split
        return records[:10]  # Limit to 10 records
    except Exception as e:
        logger.error(f"Error fetching Hugging Face dataset {dataset_name}: {str(e)}")
        return []

def fetch_vendor_logs_from_salesforce():
    """Fetch vendor logs from Salesforce with comprehensive null handling."""
    try:
        query = """
        SELECT Id, Name, Vendor__c, Work_Completion_Percentage__c, Quality_Percentage__c, 
               Incident_Severity__c, Work_Completion_Date__c, Actual_Completion_Date__c, 
               Delay_Days__c, Project__c
        FROM Vendor_Log__c
        """
        result = sf.query_all(query)
        logs = []
        for record in result['records']:
            # Skip records with missing critical fields
            if not record['Vendor__c'] or not record['Id']:
                logger.warning(f"Skipping Vendor_Log__c record with ID {record.get('Id', 'Unknown')} due to missing Vendor__c or Id")
                continue
            # Handle null values for all fields
            delay_days = record.get('Delay_Days__c', 0)
            if delay_days is None:
                logger.info(f"Delay_Days__c is null for record ID {record['Id']}, defaulting to 0")
                delay_days = 0
            work_completion = record.get('Work_Completion_Percentage__c', 0.0)
            if work_completion is None:
                logger.info(f"Work_Completion_Percentage__c is null for record ID {record['Id']}, defaulting to 0.0")
                work_completion = 0.0
            quality_percentage = record.get('Quality_Percentage__c', 0.0)
            if quality_percentage is None:
                logger.info(f"Quality_Percentage__c is null for record ID {record['Id']}, defaulting to 0.0")
                quality_percentage = 0.0
            incident_severity = record.get('Incident_Severity__c', 'None')
            if incident_severity is None:
                logger.info(f"Incident_Severity__c is null for record ID {record['Id']}, defaulting to 'None'")
                incident_severity = 'None'
            work_completion_date = record.get('Work_Completion_Date__c', 'N/A')
            if work_completion_date is None:
                logger.info(f"Work_Completion_Date__c is null for record ID {record['Id']}, defaulting to 'N/A'")
                work_completion_date = 'N/A'
            actual_completion_date = record.get('Actual_Completion_Date__c', 'N/A')
            if actual_completion_date is None:
                logger.info(f"Actual_Completion_Date__c is null for record ID {record['Id']}, defaulting to 'N/A'")
                actual_completion_date = 'N/A'
            project = record.get('Project__c', 'Unknown')
            if project is None:
                logger.info(f"Project__c is null for record ID {record['Id']}, defaulting to 'Unknown'")
                project = 'Unknown'
            name = record.get('Name', 'Unknown')
            if name is None:
                logger.info(f"Name is null for record ID {record['Id']}, defaulting to 'Unknown'")
                name = 'Unknown'
            log = VendorLog(
                vendorLogId=record['Id'],
                vendorId=name,
                vendorRecordId=record['Vendor__c'],
                workDetails=str(work_completion),
                qualityReport=str(quality_percentage),
                incidentLog=incident_severity,
                workCompletionDate=work_completion_date,
                actualCompletionDate=actual_completion_date,
                vendorLogName=name,
                delayDays=int(delay_days),
                project=project
            )
            logs.append(log)
        logger.info(f"Fetched {len(logs)} vendor logs")
        return logs
    except Exception as e:
        logger.error(f"Error fetching vendor logs from Salesforce: {str(e)}")
        return []  # Return empty list to prevent dashboard crash

def calculate_scores(log: VendorLog):
    """Calculate vendor performance scores."""
    try:
        work_completion_percentage = float(log.workDetails or 0.0)
        quality_percentage = float(log.qualityReport or 0.0)

        quality_score = quality_percentage
        timeliness_score = 100.0 if log.delayDays <= 0 else 80.0 if log.delayDays <= 3 else 60.0 if log.delayDays <= 7 else 40.0
        severity_map = {'None': 100.0, 'Low': 80.0, 'Minor': 80.0, 'Medium': 50.0, 'High': 20.0}
        safety_score = severity_map.get(log.incidentLog, 100.0)
        communication_score = (quality_score * 0.33 + timeliness_score * 0.33 + safety_score * 0.33)

        return {
            'qualityScore': round(quality_score, 2),
            'timelinessScore': round(timeliness_score, 2),
            'safetyScore': round(safety_score, 2),
            'communicationScore': round(communication_score, 2)
        }
    except Exception as e:
        logger.error(f"Error calculating scores for log {log.vendorLogId}: {str(e)}")
        return {
            'qualityScore': 0.0,
            'timelinessScore': 0.0,
            'safetyScore': 0.0,
            'communicationScore': 0.0
        }

def get_feedback(score: float, metric: str) -> str:
    """Generate feedback based on score and metric."""
    try:
        if score >= 90:
            return "Excellent: Maintain this standard"
        elif score >= 70:
            return "Good: Keep up the good work"
        elif score >= 50:
            if metric == 'Timeliness':
                return "Needs Improvement: Maintain schedules to complete tasks on time"
            elif metric == 'Safety':
                return "Needs Improvement: Implement stricter safety protocols"
            elif metric == 'Quality':
                return "Needs Improvement: Focus on improving work quality"
            else:
                return "Needs Improvement: Enhance coordination with project teams"
        else:
            if metric == 'Timeliness':
                return "Poor: Significant delays detected"
            elif metric == 'Safety':
                return "Poor: Critical safety issues identified"
            elif metric == 'Quality':
                return "Poor: Quality standards not met"
            else:
                return "Poor: Communication issues detected"
    except Exception as e:
        logger.error(f"Error generating feedback: {str(e)}")
        return "Feedback unavailable"

def generate_pdf(vendor_id: str, vendor_log_name: str, scores: dict):
    """Generate a PDF report for vendor performance."""
    try:
        filename = f'report_{vendor_id}_{datetime.now().strftime("%Y%m%d%H%M%S")}.pdf'
        c = canvas.Canvas(filename, pagesize=letter)
        c.setFont('Helvetica', 12)
        c.drawString(100, 750, 'Subcontractor Performance Report')
        c.drawString(100, 730, f'Vendor ID: {vendor_id}')
        c.drawString(100, 710, f'Vendor Log Name: {vendor_log_name}')
        c.drawString(100, 690, f'Quality Score: {scores["qualityScore"]}% ({get_feedback(scores["qualityScore"], "Quality")})')
        c.drawString(100, 670, f'Timeliness Score: {scores["timelinessScore"]}% ({get_feedback(scores["timelinessScore"], "Timeliness")})')
        c.drawString(100, 650, f'Safety Score: {scores["safetyScore"]}% ({get_feedback(scores["safetyScore"], "Safety")})')
        c.drawString(100, 630, f'Communication Score: {scores["communicationScore"]}% ({get_feedback(scores["communicationScore"], "Communication")})')
        c.save()

        with open(filename, 'rb') as f:
            pdf_content = f.read()
        os.remove(filename)
        return pdf_content
    except Exception as e:
        logger.error(f"Error generating PDF for vendor {vendor_id}: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error generating PDF: {str(e)}")

def determine_alert_flag(scores: dict, all_logs: list):
    """Determine if an alert flag should be set."""
    try:
        if not all_logs:
            return False
        avg_score = sum(scores.values()) / 4
        if avg_score < 50:
            return True
        lowest_avg = min([sum(log['scores'].values()) / 4 for log in all_logs], default=avg_score)
        return avg_score == lowest_avg
    except Exception as e:
        logger.error(f"Error determining alert flag: {str(e)}")
        return False

def store_scores_in_salesforce(log: VendorLog, scores: dict, pdf_content: bytes, alert_flag: bool):
    """Store scores and PDF in Salesforce."""
    try:
        score_record = sf.Subcontractor_Performance_Score__c.create({
            'Vendor_Log__c': log.vendorLogId,
            'Vendor__c': log.vendorRecordId,
            'Quality_Score__c': scores['qualityScore'],
            'Timeliness_Score__c': scores['timelinessScore'],
            'Safety_Score__c': scores['safetyScore'],
            'Communication_Score__c': scores['communicationScore'],
            'Alert_Flag__c': alert_flag
        })
        score_record_id = score_record['id']
        logger.info(f"Created Subcontractor_Performance_Score__c record with ID: {score_record_id}")

        pdf_base64 = base64.b64encode(pdf_content).decode('utf-8')
        content_version = sf.ContentVersion.create({
            'Title': f'Performance_Report_{log.vendorId}',
            'PathOnClient': f'report_{log.vendorId}.pdf',
            'VersionData': pdf_base64,
            'FirstPublishLocationId': score_record_id
        })
        content_version_id = content_version['id']
        content_version_record = sf.query(f"SELECT ContentDocumentId FROM ContentVersion WHERE Id = '{content_version_id}'")
        if content_version_record['totalSize'] == 0:
            logger.error(f"No ContentVersion for ID: {content_version_id}")
            raise ValueError("Failed to retrieve ContentDocumentId")
        content_document_id = content_version_record['records'][0]['ContentDocumentId']
        
        if has_pdf_link:
            pdf_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/document/download/{content_document_id}"
            sf.Subcontractor_Performance_Score__c.update(score_record_id, {'PDF_Link__c': pdf_url})
            logger.info(f"Updated Subcontractor_Performance_Score__c record with PDF URL: {pdf_url}")
        else:
            logger.info(f"Skipping PDF_Link__c update as field is not available")
    except Exception as e:
        logger.error(f"Error storing scores in Salesforce for log {log.vendorLogId}: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error storing scores: {str(e)}")

@app.post('/score')
async def score_vendor(log: VendorLog, authorization: str = Header(...)):
    """Score a vendor and generate a PDF report."""
    try:
        logger.info(f"Received Vendor Log: {log}")
        if authorization != f'Bearer {API_KEY}':
            raise HTTPException(status_code=401, detail='Invalid API key')

        scores = calculate_scores(log)
        pdf_content = generate_pdf(log.vendorId, log.vendorLogName, scores)
        pdf_base64 = base64.b64encode(pdf_content).decode('utf-8')
        alert_flag = determine_alert_flag(scores, vendor_logs)
        store_scores_in_salesforce(log, scores, pdf_content, alert_flag)

        vendor_logs.append({
            'vendorLogId': log.vendorLogId,
            'vendorId': log.vendorId,
            'vendorLogName': log.vendorLogName,
            'workDetails': log.workDetails,
            'qualityReport': log.qualityReport,
            'incidentLog': log.incidentLog,
            'workCompletionDate': log.workCompletionDate,
            'actualCompletionDate': log.actualCompletionDate,
            'delayDays': log.delayDays,
            'project': log.project,
            'scores': scores,
            'extracted': True
        })

        return {
            'vendorLogId': log.vendorLogId,
            'vendorId': log.vendorId,
            'vendorLogName': log.vendorLogName,
            'qualityScore': scores['qualityScore'],
            'timelinessScore': scores['timelinessScore'],
            'safetyScore': scores['safetyScore'],
            'communicationScore': scores['communicationScore'],
            'pdfContent': pdf_base64,
            'alert': alert_flag
        }
    except HTTPException as e:
        raise
    except Exception as e:
        logger.error(f"Error in /score endpoint: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing vendor log: {str(e)}")

@app.get('/', response_class=HTMLResponse)
async def get_dashboard():
    """Render the dashboard with vendor logs and scores."""
    try:
        global vendor_logs
        fetched_logs = fetch_vendor_logs_from_salesforce()
        for log in fetched_logs:
            if not any(existing_log['vendorLogId'] == log.vendorLogId for existing_log in vendor_logs):
                scores = calculate_scores(log)
                pdf_content = generate_pdf(log.vendorId, log.vendorLogName, scores)
                alert_flag = determine_alert_flag(scores, vendor_logs)
                store_scores_in_salesforce(log, scores, pdf_content, alert_flag)
                vendor_logs.append({
                    'vendorLogId': log.vendorLogId,
                    'vendorId': log.vendorId,
                    'vendorLogName': log.vendorLogName,
                    'workDetails': log.workDetails,
                    'qualityReport': log.qualityReport,
                    'incidentLog': log.incidentLog,
                    'workCompletionDate': log.workCompletionDate,
                    'actualCompletionDate': log.actualCompletionDate,
                    'delayDays': log.delayDays,
                    'project': log.project,
                    'scores': scores,
                    'extracted': True
                })

        html_content = """
        <html>
            <head>
                <title>Subcontractor Performance Score App</title>
                <style>
                    body { font-family: Arial, sans-serif; margin: 20px; }
                    table { width: 100%; border-collapse: collapse; margin-top: 20px; }
                    th, td { border: 1px solid #ddd; padding: 8px; text-align: left; }
                    th { background-color: #f2f2f2; }
                    h1, h2 { text-align: center; }
                    .generate-btn { 
                        display: block; 
                        margin: 20px auto; 
                        padding: 10px 20px; 
                        background-color: #4CAF50; 
                        color: white; 
                        border: none; 
                        border-radius: 5px; 
                        cursor: pointer; 
                        font-size: 16px; 
                    }
                    .generate-btn:hover { background-color: #45a049; }
                </style>
                <script>
                    async function generateScores() {
                        try {
                            const response = await fetch('/generate', { method: 'POST' });
                            if (response.ok) {
                                window.location.reload();
                            } else {
                                alert('Error generating scores');
                            }
                        } catch (error) {
                            alert('Error: ' + error.message);
                        }
                    }
                </script>
            </head>
            <body>
                <h1>SUBCONTRACTOR PERFORMANCE SCORE APP GENERATOR</h1>
                <h2>VENDOR LOGS SUBMISSION</h2>
                <table>
                    <tr>
                        <th>Vendor ID</th>
                        <th>Vendor Log Name</th>
                        <th>Project</th>
                        <th>Work Completion Percentage</th>
                        <th>Quality Percentage</th>
                        <th>Incident Severity</th>
                        <th>Work Completion Date</th>
                        <th>Actual Completion Date</th>
                        <th>Delay Days</th>
                    </tr>
        """

        if not vendor_logs:
            html_content += """
                    <tr>
                        <td colspan="9" style="text-align: center;">No vendor logs available</td>
                    </tr>
            """
        else:
            for log in vendor_logs:
                html_content += f"""
                        <tr>
                            <td>{log['vendorId']}</td>
                            <td>{log['vendorLogName']}</td>
                            <td>{log['project']}</td>
                            <td>{log['workDetails']}</td>
                            <td>{log['qualityReport']}</td>
                            <td>{log['incidentLog']}</td>
                            <td>{log['workCompletionDate']}</td>
                            <td>{log['actualCompletionDate']}</td>
                            <td>{log['delayDays']}</td>
                        </tr>
                """

        html_content += """
                </table>
                <button class="generate-btn" onclick="generateScores()">Generate</button>
                <h2>SUBCONTRACTOR PERFORMANCE SCORES</h2>
                <table>
                    <tr>
                        <th>Vendor ID</th>
                        <th>Vendor Log Name</th>
                        <th>Project</th>
                        <th>Quality Score</th>
                        <th>Timeliness Score</th>
                        <th>Safety Score</th>
                        <th>Communication Score</th>
                        <th>Alert Flag</th>
                    </tr>
        """

        if not vendor_logs:
            html_content += """
                    <tr>
                        <td colspan="8" style="text-align: center;">No scores available</td>
                    </tr>
            """
        else:
            for log in vendor_logs:
                scores = log['scores']
                alert_flag = determine_alert_flag(scores, vendor_logs)
                html_content += f"""
                        <tr>
                            <td>{log['vendorId']}</td>
                            <td>{log['vendorLogName']}</td>
                            <td>{log['project']}</td>
                            <td>{scores['qualityScore']}%</td>
                            <td>{scores['timelinessScore']}%</td>
                            <td>{scores['safetyScore']}%</td>
                            <td>{scores['communicationScore']}%</td>
                            <td>{'Checked' if alert_flag else 'Unchecked'}</td>
                        </tr>
                """

        html_content += """
                </table>
            </body>
        </html>
        """
        return HTMLResponse(content=html_content)
    except Exception as e:
        logger.error(f"Error in / endpoint: {str(e)}")
        html_content = """
        <html>
            <head><title>Error</title></head>
            <body>
                <h1>Error Loading Dashboard</h1>
                <p>An error occurred while loading the dashboard. Please check the server logs.</p>
            </body>
        </html>
        """
        return HTMLResponse(content=html_content)

@app.post('/generate')
async def generate_scores():
    """Generate scores for all vendor logs."""
    try:
        global vendor_logs
        fetched_logs = fetch_vendor_logs_from_salesforce()
        vendor_logs = []
        for log in fetched_logs:
            scores = calculate_scores(log)
            pdf_content = generate_pdf(log.vendorId, log.vendorLogName, scores)
            alert_flag = determine_alert_flag(scores, vendor_logs)
            store_scores_in_salesforce(log, scores, pdf_content, alert_flag)
            vendor_logs.append({
                'vendorLogId': log.vendorLogId,
                'vendorId': log.vendorId,
                'vendorLogName': log.vendorLogName,
                'workDetails': log.workDetails,
                'qualityReport': log.qualityReport,
                'incidentLog': log.incidentLog,
                'workCompletionDate': log.workCompletionDate,
                'actualCompletionDate': log.actualCompletionDate,
                'delayDays': log.delayDays,
                'project': log.project,
                'scores': scores,
                'extracted': True
            })
        logger.info(f"Generated scores for {len(vendor_logs)} logs")
        return {"status": "success"}
    except Exception as e:
        logger.error(f"Error in /generate endpoint: {str(e)}")
        return {"status": "error", "message": str(e)}

@app.get('/huggingface-records')
async def get_huggingface_records():
    """Fetch and return Hugging Face dataset records."""
    try:
        records = fetch_huggingface_records()
        if not records:
            logger.warning("No records fetched from Hugging Face")
            return {"records": [], "message": "No records available"}
        logger.info(f"Fetched {len(records)} Hugging Face records")
        return {"records": records}
    except Exception as e:
        logger.error(f"Error fetching Hugging Face records: {str(e)}")
        return {"records": [], "message": f"Failed to fetch Hugging Face records: {str(e)}"}

@app.get('/debug')
async def debug_info():
    """Return debug information about Salesforce and Hugging Face."""
    try:
        log_count = sf.query("SELECT COUNT() FROM Vendor_Log__c")['totalSize']
        fields = [f['name'] for f in sf.Vendor_Log__c.describe()['fields']]
        score_fields = [f['name'] for f in sf.Subcontractor_Performance_Score__c.describe()['fields']]
        hf_records = fetch_huggingface_records()
        return {
            "salesforce_connected": True,
            "vendor_log_count": log_count,
            "vendor_log_fields": fields,
            "score_fields": score_fields,
            "huggingface_records_sample": hf_records[:2]  # Limit to 2 for brevity
        }
    except Exception as e:
        logger.error(f"Debug error: {str(e)}")
        return {"salesforce_connected": False, "error": str(e)}

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