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 shutil
import base64
import pytz
import hashlib
import statistics

# 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
VALID_MILESTONES = ["Planning", "Foundation", "Walls Erected", "Interior Furnishing", "Completed"]

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

# Mock Hugging Face image analysis (replace with real model integration)
def analyze_image_with_hf(image_path):
    img = Image.open(image_path)
    detected_elements = []
    
    width, height = img.size
    total_pixels = width * height
    img_data = list(img.convert('RGB').getdata())
    gray_pixels = sum(1 for pixel in img_data if sum(pixel[:3]) / 3 < 128) / len(img_data)
    color_variance = statistics.variance([sum(pixel[:3]) / 3 for pixel in img_data]) if len(img_data) > 1 else 0
    
    # Adjusted logic to detect completed construction, overriding small size
    if (color_variance > 2000 and gray_pixels < 0.4 and any(keyword in os.path.basename(image_path).lower() for keyword in ["window", "door", "roof"])) or \
       (total_pixels > 800000 and gray_pixels < 0.4 and color_variance > 2000):
        detected_elements.extend(["roof", "windows", "doors", "cladding", "finished"])
    elif total_pixels < 200000 and color_variance < 800:
        detected_elements.append("planning")
    elif gray_pixels > 0.6 and total_pixels > 400000:
        detected_elements.append("foundation")
    elif gray_pixels > 0.35 and total_pixels > 800000 and color_variance > 1500:
        detected_elements.append("walls")
    elif gray_pixels < 0.25 or any(keyword in os.path.basename(image_path).lower() for keyword in ["interior", "plumbing", "electrical", "ceiling"]):
        detected_elements.extend(["interior", "electrical_wiring", "plumbing_pipes"])
    else:
        detected_elements.append("walls")

    return detected_elements

# Image processing and Salesforce upload with live preview
def process_image(images, project_name):
    try:
        if not images or len(images) == 0:
            return "Error: Please upload at least one image to proceed.", None, "Pending", "", "", 0

        results_html = []
        upload_statuses = []
        milestones = []
        progresses = []

        for image in images:
            img = None
            try:
                img = Image.open(image)
            except Exception as e:
                if "cannot identify image file" in str(e):
                    return f"Error: Unsupported image format '{os.path.basename(image)}'. Please use JPG, JPEG, PNG, or convert to a supported format.", None, "Failure", "", "", 0
                return f"Error: Failed to open image '{os.path.basename(image)}' - {str(e)}", None, "Failure", "", "", 0

            image_size_mb = os.path.getsize(image) / (1024 * 1024)
            if image_size_mb > 20:
                return "Error: One or more images exceed 20MB.", None, "Failure", "", "", 0
            if not str(image).lower().endswith(('.jpg', '.jpeg', '.png', '.avif')):
                return "Error: Only JPG, JPEG, PNG, or AVIF images are supported.", None, "Failure", "", "", 0

            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)

            with open(saved_image_path, 'rb') as image_file:
                image_data = base64.b64encode(image_file.read()).decode('utf-8')

            file_url = ""
            try:
                content_version = {
                    'Title': image_filename,
                    'PathOnClient': saved_image_path,
                    'VersionData': image_data
                }
                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:
                if "STORAGE_LIMIT_EXCEEDED" in str(e):
                    file_url = "Upload skipped due to storage limit"
                else:
                    return f"Error: Failed to upload image to Salesforce - {str(e)}", None, "Failure", "", "", 0

            detected_elements = analyze_image_with_hf(image)

            completed_tasks = []
            not_completed_tasks = []
            all_tasks = {
                "Planning": [
                    "Initial project outline and objectives have been established.",
                    "Preliminary designs and architectural plans are drafted.",
                    "Stakeholder meetings and initial approvals are completed."
                ],
                "Foundation": [
                    "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."
                ],
                "Walls Erected": [
                    "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."
                ],
                "Interior Furnishing": [
                    "Plumbing and electrical groundwork installations are completed.",
                    "Interior walls are fully implemented.",
                    "Flooring installation is completed.",
                    "Painting and fixtures are installed."
                ],
                "Completed": [
                    "Roofing installation and weatherproofing are completed.",
                    "Windows, doors, and exterior cladding are installed.",
                    "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."
                ]
            }

            if any(elem == "planning" for elem in detected_elements):
                completed_tasks.extend(all_tasks["Planning"])
                final_milestone = "Planning"
                percent_complete = 10
            elif any(elem == "foundation" for elem in detected_elements):
                completed_tasks.extend(all_tasks["Foundation"])
                final_milestone = "Foundation"
                percent_complete = 30
            elif any(elem == "walls" for elem in detected_elements):
                completed_tasks.extend(all_tasks["Walls Erected"])
                final_milestone = "Walls Erected"
                percent_complete = 50
            elif any(elem in ["interior", "electrical_wiring", "plumbing_pipes"] for elem in detected_elements):
                completed_tasks.extend(all_tasks["Interior Furnishing"])
                final_milestone = "Interior Furnishing"
                percent_complete = 80
            elif all(elem in ["roof", "windows", "doors", "cladding", "finished"] for elem in detected_elements if elem in ["roof", "windows", "doors", "cladding", "finished"]):
                completed_tasks.extend(all_tasks["Completed"])
                final_milestone = "Completed"
                percent_complete = 100
            else:
                completed_tasks.extend(all_tasks["Walls Erected"])
                final_milestone = "Walls Erected"
                percent_complete = 50

            all_possible_tasks = [task for sublist in all_tasks.values() for task in sublist]
            not_completed_tasks = [task for task in all_possible_tasks if task not in completed_tasks]

            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])

            if final_milestone == "Completed":
                result_html = f"""
                <div style="font-family: Arial, sans-serif; padding: 20px; background-color: #f9f9f9; border-radius: 10px; margin-bottom: 20px;">
                    <h2 style="color: #2c3e50; text-align: center;">Project Summary - {os.path.basename(image)}</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;">Milestone Timeline</h3>
                    <div style="display: flex; justify-content: space-between; align-items: center; margin-bottom: 20px;">
                        <span style="color: #bdc3c7;">Planning</span>
                        <span style="color: #bdc3c7;">Foundation</span>
                        <span style="color: #bdc3c7;">Walls Erected</span>
                        <span style="color: #bdc3c7;">Interior Furnishing</span>
                        <span style="color: #2ecc71;">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>
                    <p style="color: green;">Project is fully completed as of 02:35 PM IST, June 20, 2025.</p>
                </div>
                """
            else:
                result_html = f"""
                <div style="font-family: Arial, sans-serif; padding: 20px; background-color: #f9f9f9; border-radius: 10px; margin-bottom: 20px;">
                    <h2 style="color: #2c3e50; text-align: center;">Project Summary - {os.path.basename(image)}</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;">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 == 'Interior Furnishing' else '#bdc3c7'};">Interior Furnishing</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>
                    <p style="color: orange;">Construction is in progress at {final_milestone} stage as of 02:35 PM IST, June 20, 2025.</p>
                </div>
                """

            results_html.append(result_html)
            upload_statuses.append("Success" if file_url != "Upload skipped due to storage limit" else "Pending")
            milestones.append(final_milestone)
            progresses.append(f"{percent_complete}%")

        combined_html = "<div>" + "".join(results_html) + "</div>"

        # Set current time to 02:35 PM IST, June 20, 2025
        now = local_timezone.localize(datetime(2025, 6, 20, 14, 35))
        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": milestones[0] if milestones else "Pending",
            "Completion_Percentage__c": int(progresses[0].rstrip('%')) if progresses else 0,
            "Last_Updated_On__c": local_time,
            "Upload_Status__c": upload_statuses[0] if upload_statuses else "Pending",
            "Comments__c": milestones[0] if milestones else "Pending",
            "Last_Updated_Image__c": file_url if file_url else ""
        }

        try:
            if file_url and file_url != "Upload skipped due to storage limit":
                sf.Construction__c.create(record)
        except Exception as e:
            return f"Error: Failed to update Salesforce - {str(e)}", None, "Failure", "", "", 0

        return combined_html, Image.open(images[0]) if images else None, ",".join(upload_statuses), ",".join(milestones), "", ",".join(progresses)

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

# Gradio UI with enhanced styling and live image preview
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;
    }
    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>")
    with gr.Row():
        image_input = gr.Files(type="filepath", label="Upload Construction Site Photos (JPG/PNG, ≤ 20MB)")
        image_preview = gr.Image(label="Image Preview", interactive=False)
        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")
    confidence = gr.Textbox(label="Confidence Score")  # Kept for compatibility, but unused
    progress = gr.Textbox(label="Completion Percentage", interactive=False)

    # Update preview on file upload
    image_input.change(
        fn=lambda x: Image.open(x[0]) if x else None,
        inputs=[image_input],
        outputs=[image_preview]
    )
    submit_button.click(
        fn=process_image,
        inputs=[image_input, project_name_input],
        outputs=[output_html, image_preview, upload_status, milestone, confidence, progress]
    )

demo.launch(share=True)