import gradio as gr
from PIL import Image
import os
from dotenv import load_dotenv
from simple_salesforce import Salesforce
from datetime import datetime
import hashlib
import shutil
import base64
import pytz

# Load environment variables
load_dotenv()
SF_USERNAME = os.getenv("SF_USERNAME")
SF_PASSWORD = os.getenv("SF_PASSWORD")
SF_SECURITY_TOKEN = os.getenv("SF_SECURITY_TOKEN")

# Validate Salesforce credentials
if not all([SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN]):
    raise ValueError("Missing Salesforce credentials. Set SF_USERNAME, SF_PASSWORD, and SF_SECURITY_TOKEN in environment variables.")

# Initialize Salesforce connection
try:
    sf = Salesforce(
        username=SF_USERNAME,
        password=SF_PASSWORD,
        security_token=SF_SECURITY_TOKEN,
        domain='login'
    )
except Exception as e:
    print(f"Salesforce connection failed: {str(e)}")
    raise

# Valid milestones in sequential order
VALID_MILESTONES = ["Planning", "Foundation", "Walls Erected", "Completed"]
MILESTONE_WEIGHTS = {
    "Planning": 1,
    "Foundation": 2,
    "Walls Erected": 3,
    "Completed": 4
}

# Adjust the timezone to your local timezone
local_timezone = pytz.timezone("Asia/Kolkata")

# Image processing and Salesforce upload
def process_image(images, project_name):
    try:
        if not images or len(images) == 0:
            return "Error: Please upload at least one image to proceed.", "Pending", "", "", 0
        if len(images) < 2:
            return "Error: Please upload at least one indoor and one outdoor image for accurate milestone detection.", "Pending", "", "", 0

        # Process each image
        image_milestones = []
        image_types = []
        
        for image in images:
            img = Image.open(image)
            image_size_mb = os.path.getsize(image) / (1024 * 1024)
            if image_size_mb > 20:
                return "Error: One or more images exceed 20MB.", "Failure", "", "", 0
            if not str(image).lower().endswith(('.jpg', '.jpeg', '.png')):
                return "Error: Only JPG/PNG images are supported.", "Failure", "", "", 0

            # Save image to public folder temporarily before uploading to Salesforce
            upload_dir = "public_uploads"
            os.makedirs(upload_dir, exist_ok=True)
            unique_id = datetime.now().strftime("%Y%m%d%H%M%S")
            image_filename = f"{unique_id}_{os.path.basename(image)}"
            saved_image_path = os.path.join(upload_dir, image_filename)
            shutil.copy(image, saved_image_path)

            # Convert image to base64 before uploading to Salesforce
            with open(saved_image_path, 'rb') as image_file:
                image_data = base64.b64encode(image_file.read()).decode('utf-8')

            # Create the ContentVersion record in Salesforce
            content_version = {
                'Title': image_filename,
                'PathOnClient': saved_image_path,
                'VersionData': image_data
            }

            # Upload the file to Salesforce
            try:
                content_version_result = sf.ContentVersion.create(content_version)
                content_version_id = content_version_result['id']
                file_url = f"https://sathkruthatechsolutionspri8-dev-ed.develop.lightning.force.com/{content_version_id}"
            except Exception as e:
                return f"Error: Failed to upload image to Salesforce - {str(e)}", "Failure", "", "", 0

            # Classify image as indoor or outdoor based on filename
            filename_lower = os.path.basename(image).lower()
            is_indoor = any(keyword in filename_lower for keyword in ["indoor", "interior", "inside"])
            image_type = "Indoor" if is_indoor else "Outdoor"
            image_types.append(image_type)

            # Enhanced milestone detection logic
            # Use filename keywords to simulate content analysis
            milestone = "Planning"  # Default
            if any(keyword in filename_lower for keyword in ["site", "clearing", "planning", "design"]):
                milestone = "Planning"
            elif any(keyword in filename_lower for keyword in ["foundation", "footing", "slab", "excavation"]):
                milestone = "Foundation"
            elif any(keyword in filename_lower for keyword in ["wall", "structure", "beam", "column", "facade"]):
                milestone = "Walls Erected"
            elif any(keyword in filename_lower for keyword in ["electrical", "plumbing", "hvac", "finish", "completed"]):
                milestone = "Completed"

            # Adjust milestone based on image type
            if image_type == "Indoor" and milestone in ["Planning", "Foundation"]:
                milestone = "Walls Erected"  # Indoor images imply at least walls are up
            elif image_type == "Outdoor" and milestone == "Completed":
                milestone = "Walls Erected"  # Outdoor completion needs indoor confirmation

            image_milestones.append(milestone)

        # Validate and aggregate milestones
        if not any(t == "Indoor" for t in image_types) or not any(t == "Outdoor" for t in image_types):
            return "Error: Both indoor and outdoor images are required for accurate milestone detection.", "Pending", "", "", 0

        # Ensure sequential milestone logic
        max_milestone_index = max(MILESTONE_WEIGHTS[m] for m in image_milestones)
        final_milestone = [m for m, w in MILESTONE_WEIGHTS.items() if w == max_milestone_index][0]

        # If "Completed" is detected, ensure indoor images confirm it
        if final_milestone == "Completed" and not any(m == "Completed" and t == "Indoor" for m, t in zip(image_milestones, image_types)):
            final_milestone = "Walls Erected"

        milestone_completion_map = {
            "Planning": 10,
            "Foundation": 30,
            "Walls Erected": 50,
            "Completed": 100,
        }
        percent_complete = milestone_completion_map.get(final_milestone, 0)

        completion_details = {
            "Planning": {
                "completed": [
                    "Initial project outline and objectives have been established.",
                    "Preliminary designs and architectural plans are drafted.",
                    "Stakeholder meetings and initial approvals are completed."
                ],
                "not_completed": [
                    "Detailed construction plans and blueprints are pending finalization.",
                    "Permits and regulatory approvals are yet to be obtained.",
                    "Contractor selection and procurement processes are not yet complete."
                ]
            },
            "Foundation": {
                "completed": [
                    "Site preparation, including clearing and leveling, is finished.",
                    "Excavation for the foundation has been completed.",
                    "Concrete pouring for the foundation, including footings and slabs, is done.",
                    "Initial structural inspections for the foundation have been passed."
                ],
                "not_completed": [
                    "Plumbing and electrical groundwork installations are pending.",
                    "Backfilling and site grading around the foundation are not yet done.",
                    "Above-ground structural work, such as columns and walls, has not started."
                ]
            },
            "Walls Erected": {
                "completed": [
                    "The concrete framework, including columns and beams, is in place.",
                    "All structural walls have been erected and stabilized.",
                    "Temporary scaffolding and safety measures are installed for ongoing work.",
                    "Initial inspections for structural integrity have been completed."
                ],
                "not_completed": [
                    "Roofing installation and weatherproofing are pending.",
                    "Windows, doors, and exterior cladding are not yet installed.",
                    "Interior walls, electrical, and plumbing systems are still to be implemented."
                ]
            },
            "Completed": {
                "completed": [
                    "The concrete framework, including columns, beams, and floor slabs, is fully constructed.",
                    "Exterior walls, windows, and cladding are installed, completing the building's facade.",
                    "Interior work, including electrical, plumbing, and HVAC systems, is fully implemented.",
                    "Finishing touches, such as flooring, painting, and fixtures, are completed.",
                    "All phases of the project are finished, including final inspections and approvals."
                ],
                "not_completed": [
                    "There should be no more pending work as the project is fully completed."
                ]
            }
        }

        completed_tasks = completion_details[final_milestone]["completed"]
        not_completed_tasks = completion_details[final_milestone]["not_completed"]

        completed_html = "".join([f'<li style="color: green;">✔ {task}</li>' for task in completed_tasks])
        not_completed_html = "".join([f'<li style="color: red;">✘ {task}</li>' for task in not_completed_tasks])

        # Enhanced result HTML with image type feedback
        image_summary = "".join([f'<li>{os.path.basename(img)} ({img_type}): {milestone}</li>' for img, img_type, milestone in zip(images, image_types, image_milestones)])
        
        result_html = f"""
        <div style="font-family: Arial, sans-serif; padding: 20px; background-color: #f9f9f9; border-radius: 10px;">
            <h2 style="color: #2c3e50; text-align: center;">Project Summary</h2>
            <div style="display: flex; justify-content: space-around; margin-bottom: 20px;">
                <div style="text-align: center;">
                    <h3 style="color: #34495e;">Detected Milestone</h3>
                    <p style="font-size: 18px; font-weight: bold;">{final_milestone}</p>
                </div>
                <div style="text-align: center;">
                    <h3 style="color: #34495e;">Completion</h3>
                    <progress value="{percent_complete}" max="100" style="width: 200px; height: 20px;"></progress>
                    <p>{percent_complete}%</p>
                </div>
            </div>

            <h3 style="color: #2c3e50;">Image Analysis</h3>
            <ul style="padding-left: 20px; margin-bottom: 20px;">
                {image_summary}
            </ul>

            <h3 style="color: #2c3e50;">Milestone Timeline</h3>
            <div style="display: flex; justify-content: space-between; align-items: center; margin-bottom: 20px;">
                <span style="color: {'#2ecc71' if final_milestone == 'Planning' else '#bdc3c7'};">Planning</span>
                <span style="color: {'#2ecc71' if final_milestone == 'Foundation' else '#bdc3c7'};">Foundation</span>
                <span style="color: {'#2ecc71' if final_milestone == 'Walls Erected' else '#bdc3c7'};">Walls Erected</span>
                <span style="color: {'#2ecc71' if final_milestone == 'Completed' else '#bdc3c7'};">Completed</span>
            </div>

            <details style="margin-bottom: 20px;">
                <summary style="color: #2c3e50; font-weight: bold;">Completed Tasks</summary>
                <ul style="padding-left: 20px;">
                    {completed_html}
                </ul>
            </details>

            <details style="margin-bottom: 20px;">
                <summary style="color: #2c3e50; font-weight: bold;">Not Completed Tasks</summary>
                <ul style="padding-left: 20px;">
                    {not_completed_html}
                </ul>
            </details>
        </div>
        """

        now = datetime.now(local_timezone)
        local_time = now.strftime("%Y-%m-%dT%H:%M:%S") + now.strftime("%z")[:-2] + ":" + now.strftime("%z")[-2:]

        record = {
            "Name__c": project_name,
            "Current_Milestone__c": final_milestone,
            "Completion_Percentage__c": percent_complete,
            "Last_Updated_On__c": local_time,
            "Upload_Status__c": "Success",
            "Comments__c": f"{final_milestone}",
            "Last_Updated_Image__c": file_url
        }

        try:
            sf.Construction__c.create(record)
        except Exception as e:
            return f"Error: Failed to update Salesforce - {str(e)}", "Failure", "", "", 0

        return result_html, "Success", final_milestone, f"{percent_complete}%"

    except Exception as e:
        return f"Error: {str(e)}", "Failure", "", "", "0%"

# Gradio UI
with gr.Blocks(css="""
    .gradio-container {
        background-color: #f0f4f8;
        font-family: Arial, sans-serif;
    }
    .title {
        color: #2c3e50;
        font-size: 24px;
        text-align: center;
        font-weight: bold;
    }
    .gradio-row {
        text-align: center;
    }
    .gradio-container .output {
        text-align: center;
    }
    .gradio-container .input {
        text-align: center;
    }
    .gradio-container .button {
        display: block;
        margin: 0 auto;
        background-color: #3498db;
        color: white;
        border: none;
        padding: 10px 20px;
        border-radius: 5px;
        cursor: pointer;
    }
    .gradio-container .button:hover {
        background-color: #2980b9;
    }
    progress::-webkit-progress-value {
        background-color: #2ecc71;
        border-radius: 5px;
    }
    .gradio-container progress::-webkit-progress-bar {
        background-color: #ecf0f1;
        border-radius: 5px;
    }
    details summary {
        cursor: pointer;
        padding: 10px;
        background-color: #ecf0f1;
        border-radius: 5px;
    }
    details ul {
        margin-top: 10px;
    }
""") as demo:
    gr.Markdown("<h1 class='title'>Construction Progress Analyzer</h1>")
    gr.Markdown("""
        <p style='text-align: center; color: #34495e;'>
            Upload at least one indoor and one outdoor image of the construction site (JPG/PNG, ≤ 20MB each).<br>
            Use descriptive filenames (e.g., 'outdoor_foundation.jpg', 'indoor_electrical.jpg') for best results.
        </p>
    """)
    with gr.Row():
        image_input = gr.Files(type="filepath", label="Upload Construction Site Photos (JPG/PNG, ≤ 20MB)")
        project_name_input = gr.Textbox(label="Project Name (Required)", placeholder="e.g. Project_12345")

    submit_button = gr.Button("Process Image")
    output_html = gr.HTML(label="Result")
    upload_status = gr.Textbox(label="Upload Status")
    milestone = gr.Textbox(label="Detected Milestone")
    progress = gr.Textbox(label="Completion Percentage", interactive=False)

    submit_button.click(
        fn=process_image,
        inputs=[image_input, project_name_input],
        outputs=[output_html, upload_status, milestone, progress]
    )

demo.launch(share=True)