import gradio as gr
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from io import BytesIO
import os
import logging
import base64
import shutil
import tempfile
from simple_salesforce import Salesforce
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
from fastapi import FastAPI, Form, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Configure logging to show detailed messages
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Salesforce credentials (loaded from environment variables)
SALESFORCE_USERNAME = os.getenv("SALESFORCE_USERNAME")
SALESFORCE_PASSWORD = os.getenv("SALESFORCE_PASSWORD")
SALESFORCE_SECURITY_TOKEN = os.getenv("SALESFORCE_SECURITY_TOKEN")
SALESFORCE_DOMAIN = os.getenv("SALESFORCE_DOMAIN", "login")

# Validate that credentials are set
if not all([SALESFORCE_USERNAME, SALESFORCE_PASSWORD, SALESFORCE_SECURITY_TOKEN]):
    logger.error("Salesforce credentials not set in environment variables.")
    raise ValueError("Missing Salesforce credentials in environment variables.")

logger.debug("Using Salesforce credentials - Username and Security Token loaded from environment variables.")

# Function to authenticate with Salesforce
def get_salesforce_connection():
    try:
        logger.debug("Attempting to connect to Salesforce...")
        sf = Salesforce(
            username=SALESFORCE_USERNAME,
            password=SALESFORCE_PASSWORD,
            security_token=SALESFORCE_SECURITY_TOKEN,
            domain=SALESFORCE_DOMAIN
        )
        logger.info("Salesforce connection successful.")
        result = sf.query("SELECT Id FROM User LIMIT 1")
        logger.debug(f"Successfully queried Salesforce to confirm connection. Result: {result}")
        return sf
    except Exception as e:
        logger.error(f"Failed to connect to Salesforce: {str(e)}", exc_info=True)
        return None

# Function to upload a file to Salesforce as a ContentVersion
def upload_file_to_salesforce(file_path, file_name, record_id=None):
    try:
        sf = get_salesforce_connection()
        if not sf:
            logger.error("Salesforce connection failed. Cannot upload file.")
            return None

        with open(file_path, "rb") as f:
            file_data = f.read()

        encoded_file_data = base64.b64encode(file_data).decode('utf-8')
        logger.debug(f"Uploading file {file_name} for record ID: {record_id}")
        content_version_data = {
            "Title": file_name,
            "PathOnClient": file_name,
            "VersionData": encoded_file_data,
        }
        
        if record_id:
            content_version_data["FirstPublishLocationId"] = record_id

        content_version = sf.ContentVersion.create(content_version_data)
        logger.info(f"File uploaded to Salesforce with ContentVersion ID: {content_version['id']}")
        return content_version["id"]
    except Exception as e:
        logger.error(f"Error uploading file to Salesforce: {str(e)}", exc_info=True)
        return None

# Function to generate PDF
def generate_pdf(record_data):
    try:
        logger.debug("Generating PDF...")
        pdf_file = BytesIO()
        c = canvas.Canvas(pdf_file, pagesize=letter)
        c.drawString(100, 750, f"Project Title: {record_data['project_title']}")
        c.drawString(100, 730, f"Estimated Duration: {record_data['estimated_duration']} days")
        c.drawString(100, 710, f"AI Plan Score: {record_data['ai_plan_score']}%")
        c.drawString(100, 690, f"Status: {record_data['status']}")
        c.drawString(100, 670, f"Risk Tags: {record_data['risk_tags']}")
        c.save()
        pdf_file.seek(0)
        logger.debug("PDF generated successfully.")
        return pdf_file
    except Exception as e:
        logger.error(f"Error generating PDF: {str(e)}", exc_info=True)
        return None

# Function to upload PDF to Salesforce and get its URL
def upload_pdf_to_salesforce(pdf_file, project_title, record_id=None):
    try:
        sf = get_salesforce_connection()
        if not sf:
            logger.error("Salesforce connection failed. Cannot upload PDF.")
            return None, None

        encoded_pdf_data = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
        logger.debug(f"Uploading PDF for project: {project_title}, record ID: {record_id}")
        content_version_data = {
            "Title": f"{project_title} - Gantt Chart PDF",
            "PathOnClient": f"{project_title}_Gantt_Chart.pdf",
            "VersionData": encoded_pdf_data,
        }
        
        if record_id:
            content_version_data["FirstPublishLocationId"] = record_id

        content_version = sf.ContentVersion.create(content_version_data)
        content_version_id = content_version["id"]
        logger.info(f"PDF uploaded to Salesforce with ContentVersion ID: {content_version_id}")

        result = sf.query(f"SELECT Id, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version_id}'")
        if not result['records']:
            logger.error("No records returned for ContentVersion query")
            return content_version_id, None
            
        content_document_id = result['records'][0]['ContentDocumentId']
        file_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version_id}"
        logger.debug(f"Generated PDF URL: {file_url}")
        return content_version_id, file_url
    except Exception as e:
        logger.error(f"Error uploading PDF to Salesforce: {str(e)}", exc_info=True)
        return None, None

# Function to create or update project timeline in Salesforce
def send_to_salesforce(project_title, gantt_chart_url, ai_plan_score, estimated_duration, status="Draft", record_id=None, location="", weather_type="", work_items=None, work_items_id=None):
    try:
        logger.debug("Starting send_to_salesforce function...")
        sf = get_salesforce_connection()
        if not sf:
            logger.error("Salesforce connection failed. Cannot proceed with record creation/update.")
            return None

        try:
            obj_description = sf.AI_Project_Timeline__c.describe()
            logger.debug("AI_Project_Timeline__c object exists and is accessible.")
            available_fields = [field['name'] for field in obj_description['fields']]
            logger.debug(f"Available fields on AI_Project_Timeline__c: {available_fields}")
        except Exception as e:
            logger.error(f"Error: AI_Project_Timeline__c object not found or inaccessible: {str(e)}")
            return None

        sf_data = {
            "Name": project_title[:80],
            "Project_Title__c": project_title,
            "Estimated_Duration__c": estimated_duration,
            "AI_Plan_Score__c": ai_plan_score,
            "Status__c": status,
            "Location__c": location,
            "Weather_Type__c": weather_type,
        }
        
        if gantt_chart_url:
            sf_data["Gantt_Chart_PDF__c"] = gantt_chart_url
            
        if work_items_id:
            sf_data["Work_Items__c"] = work_items_id

        logger.debug(f"Prepared Salesforce data: {sf_data}")

        if record_id:
            try:
                logger.info(f"Attempting to update Salesforce record with ID: {record_id}")
                sf.AI_Project_Timeline__c.update(record_id, sf_data)
                logger.info(f"Successfully updated Salesforce record with ID: {record_id}")
                return record_id
            except Exception as e:
                logger.error(f"Error updating record {record_id}: {str(e)}")
                record_id = None

        logger.info("Creating new Salesforce record...")
        project_record = sf.AI_Project_Timeline__c.create(sf_data)
        if not project_record.get('id'):
            logger.error("Failed to create record, no ID returned")
            return None
            
        new_record_id = project_record['id']
        logger.info(f"Created new Salesforce record with ID: {new_record_id}")
        return new_record_id
    
    except Exception as e:
        logger.error(f"Error sending data to Salesforce: {str(e)}", exc_info=True)
        if hasattr(e, 'content') and e.content:
            logger.error(f"Salesforce API response: {e.content}")
        return None

# Function to generate Gantt chart
def generate_project_timeline(boq_file, weather, workforce, location, project_title):
    temp_dir = None
    try:
        logger.debug("Processing BOQ data...")
        if not boq_file:
            raise ValueError("No file uploaded")

        temp_dir = tempfile.mkdtemp()
        output_filename = f"gantt_chart_{project_title.replace(' ', '')}{id(boq_file)}.png"
        output_path = os.path.join(temp_dir, output_filename)
        logger.debug(f"Gantt chart will be saved to: {output_path}")

        if isinstance(boq_file, str):
            df = pd.read_csv(boq_file)
        else:
            df = pd.read_csv(boq_file.name)

        if "Task Name" not in df.columns or "Duration" not in df.columns:
            raise ValueError("CSV must contain 'Task Name' and 'Duration' columns")

        task_names = df["Task Name"].tolist()
        task_durations = df["Duration"].tolist()
        logger.debug(f"Tasks: {task_names}, Durations: {task_durations}")

        fig, ax = plt.subplots(figsize=(10, 5))
        ax.barh(task_names, task_durations, color="skyblue")
        ax.set_xlabel("Duration (days)")
        ax.set_title("Project Timeline Gantt Chart")
        fig.savefig(output_path, format="png", bbox_inches="tight")
        plt.close(fig)

        risk_tags = [
            f"{task} - {'High' if weather == 'rainy' and duration > 5 else 'Low'} Risk (Weather)"
            for task, duration in zip(task_names, task_durations)
        ]
        risk_tags_str = "\n".join(risk_tags)
        logger.debug(f"Generated risk tags: {risk_tags_str}")
        logger.info("Gantt chart and risk tags generated successfully.")
        return output_path, risk_tags_str, temp_dir
    except Exception as e:
        logger.error(f"Error generating project timeline: {str(e)}", exc_info=True)
        if temp_dir and os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
        return None, str(e), None

# Gradio interface function
def gradio_interface(boq_file, weather, workforce, location, project_title):
    temp_dir = None
    try:
        logger.info("Starting gradio_interface...")
        if not boq_file:
            return None, "Error: No BOQ file uploaded"

        boq_file_path = boq_file.name if hasattr(boq_file, 'name') else boq_file
        file_path, risk_tags, temp_dir = generate_project_timeline(boq_file_path, weather, workforce, location, project_title)
        if not file_path:
            return None, f"Error: Failed to generate timeline: {risk_tags}"

        df = pd.read_csv(boq_file_path)
        estimated_duration = sum(df["Duration"])
        ai_plan_score = min(100, max(0, 100 - (estimated_duration / 100)))
        logger.debug(f"Estimated duration: {estimated_duration}, AI plan score: {ai_plan_score}")

        record_id = send_to_salesforce(
            project_title=project_title,
            gantt_chart_url="",
            ai_plan_score=ai_plan_score,
            estimated_duration=estimated_duration,
            status="Draft",
            record_id=None,
            location=location,
            weather_type=weather
        )
        
        if not record_id:
            return None, f"Error: Failed to create Salesforce record - check logs for details\n\nRisk Tags:\n{risk_tags}"

        work_items_id = upload_file_to_salesforce(boq_file_path, "Boq_data.csv", record_id)
        if not work_items_id:
            logger.warning("Failed to upload BOQ file, but proceeding with record creation")

        record_data = {
            "project_title": project_title,
            "estimated_duration": estimated_duration,
            "ai_plan_score": ai_plan_score,
            "status": "Draft",
            "risk_tags": risk_tags,
        }
        pdf_file = generate_pdf(record_data)
        if not pdf_file:
            logger.warning("Failed to generate PDF, but proceeding with record creation")

        pdf_content_id, pdf_url = None, None
        if pdf_file:
            pdf_content_id, pdf_url = upload_pdf_to_salesforce(pdf_file, project_title, record_id)
            if not pdf_content_id:
                logger.warning("Failed to upload PDF, but proceeding with record creation")

        update_result = send_to_salesforce(
            project_title=project_title,
            gantt_chart_url=pdf_url if pdf_url else "",
            ai_plan_score=ai_plan_score,
            estimated_duration=estimated_duration,
            status="Draft",
            record_id=record_id,
            location=location,
            weather_type=weather,
            work_items_id=work_items_id if work_items_id else ""
        )
        if not update_result:
            logger.warning("Failed to update record with PDF URL, but record was created")

        image_content_id = upload_file_to_salesforce(fileLf_path, f"{project_title}_Gantt_Chart.png", record_id)
        image_url = None
        if image_content_id:
            sf = get_salesforce_connection()
            if sf:
                image_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{image_content_id}"
                logger.debug(f"Generated image URL: {image_url}")

        logger.info("Gradio interface completed successfully.")
        return image_url if image_url else file_path, f"Successfully created Salesforce record ID: {record_id}\n\nRisk Tags:\n{risk_tags}"
    except Exception as e:
        logger.error(f"Error in Gradio interface: {str(e)}", exc_info=True)
        return None, f"Error in Gradio interface: {str(e)}"
    finally:
        if temp_dir and os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
            logger.debug(f"Cleaned up temporary directory: {temp_dir}")

# Create Gradio interface
demo = gr.Blocks(theme="default")
with demo:
    gr.Markdown("## AI Civil Work Planner")
    gr.Markdown("Generate a project timeline (Gantt chart) and risk tags based on BOQ data and site parameters.")

    with gr.Row():
        with gr.Column():
            boq_file = gr.File(label="Upload BOQ Data (CSV format)")
            weather = gr.Dropdown(label="Weather", choices=["sunny", "rainy", "cloudy"], value="sunny")
            workforce = gr.Number(label="Workforce Size", value=10, precision=0)
            location = gr.Textbox(label="Location", placeholder="Enter project location")
            project_title = gr.Textbox(label="Project Title", placeholder="Enter project title")
            submit_btn = gr.Button("Generate Timeline")

        with gr.Column():
            output_image = gr.Image(label="Gantt Chart")
            risk_tags = gr.Textbox(label="Risk Tags and Salesforce Status")

    submit_btn.click(
        fn=gradio_interface,
        inputs=[boq_file, weather, workforce, location, project_title],
        outputs=[output_image, risk_tags],
    )

# Create a FastAPI app with CORS support
app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://aiplannerforcivilworktimel2-dev-ed.develop.lightning.force.com"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Mount directory for temporary files (e.g., Gantt chart PNGs)
app.mount("/static", StaticFiles(directory=tempfile.gettempdir()), name="static")

# Health check endpoint to verify server status
@app.get("/health")
async def health_check():
    return {"status": "healthy"}

# FastAPI endpoint for processing BOQ files and interacting with Salesforce
@app.post("/api/gradio_interface")
async def api_gradio_interface(
    boq_file: UploadFile = File(...),
    weather: str = Form(...),
    workforce: int = Form(...),
    location: str = Form(...),
    project_title: str = Form(...)
):
    temp_dir = None
    try:
        logger.info("Starting api_gradio_interface...")

        temp_dir = tempfile.mkdtemp()
        boq_file_path = os.path.join(temp_dir, boq_file.filename)
        with open(boq_file_path, "wb") as f:
            f.write(boq_file.file.read())

        file_path, risk_tags, temp_dir = generate_project_timeline(boq_file_path, weather, workforce, location, project_title)
        if not file_path:
            return JSONResponse({"error": f"Failed to generate timeline: {risk_tags}"}, status_code=400)

        df = pd.read_csv(boq_file_path)
        estimated_duration = sum(df["Duration"])
        ai_plan_score = min(100, max(0, 100 - (estimated_duration / 100)))
        logger.debug(f"Estimated duration: {estimated_duration}, AI plan score: {ai_plan_score}")

        record_id = send_to_salesforce(
            project_title=project_title,
            gantt_chart_url="",
            ai_plan_score=ai_plan_score,
            estimated_duration=estimated_duration,
            status="Draft",
            record_id=None,
            location=location,
            weather_type=weather
        )
        
        if not record_id:
            return JSONResponse({
                "error": f"Failed to create Salesforce record - check logs for details",
                "text": f"Risk Tags:\n{risk_tags}"
            }, status_code=500)

        work_items_id = upload_file_to_salesforce(boq_file_path, "Boq_data.csv", record_id)
        if not work_items_id:
            logger.warning("Failed to upload BOQ file, but proceeding with record creation")

        record_data = {
            "project_title": project_title,
            "estimated_duration": estimated_duration,
            "ai_plan_score": ai_plan_score,
            "status": "Draft",
            "risk_tags": risk_tags,
        }
        pdf_file = generate_pdf(record_data)
        if not pdf_file:
            logger.warning("Failed to generate PDF, but proceeding with record creation")

        pdf_content_id, pdf_url = None, None
        if pdf_file:
            pdf_content_id, pdf_url = upload_pdf_to_salesforce(pdf_file, project_title, record_id)
            if not pdf_content_id:
                logger.warning("Failed to upload PDF, but proceeding with record creation")

        update_result = send_to_salesforce(
            project_title=project_title,
            gantt_chart_url=pdf_url if pdf_url else "",
            ai_plan_score=ai_plan_score,
            estimated_duration=estimated_duration,
            status="Draft",
            record_id=record_id,
            location=location,
            weather_type=weather,
            work_items_id=work_items_id if work_items_id else ""
        )
        if not update_result:
            logger.warning("Failed to update record with PDF URL, but record was created")

        image_content_id = upload_file_to_salesforce(file_path, f"{project_title}_Gantt_Chart.png", record_id)
        image_url = None
        if image_content_id:
            sf = get_salesforce_connection()
            if sf:
                image_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{image_content_id}"
                logger.debug(f"Generated image URL: {image_url}")

        logger.info("API gradio interface completed successfully.")
        return JSONResponse({
            "image": image_url if image_url else f"/static/{os.path.basename(file_path)}",
            "text": f"Successfully created Salesforce record ID: {record_id}\n\nRisk Tags:\n{risk_tags}"
        })
    except Exception as e:
        logger.error(f"Error in API gradio interface: {str(e)}", exc_info=True)
        return JSONResponse({"error": f"Error in API gradio interface: {str(e)}"}, status_code=500)
    finally:
        if temp_dir and os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
            logger.debug(f"Cleaned up temporary directory: {temp_dir}")

if __name__ == "__main__":
    # Run Gradio UI
    demo.launch(server_name="0.0.0.0", server_port=7860)