import gradio as gr
from PIL import Image
import numpy as np
import os
from dotenv import load_dotenv
from simple_salesforce import Salesforce
from datetime import datetime
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")

if not all([SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN]):
    raise ValueError("Missing Salesforce credentials.")

# Connect to Salesforce
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 = ["Planning", "Foundation", "Walls Erected", "Completed"]
local_timezone = pytz.timezone("Asia/Kolkata")

# Milestone detection based on image brightness
def detect_milestone_from_image(image_path):
    img = Image.open(image_path).convert("L")
    img_array = np.array(img)
    brightness = np.mean(img_array)

    if brightness < 80:
        return "Planning"
    elif 80 <= brightness < 130:
        return "Foundation"
    elif 130 <= brightness < 180:
        return "Walls Erected"
    else:
        return "Completed"

def process_image(images, project_name):
    try:
        # Ensure images is a list and not empty
        if not images or len(images) == 0 or not isinstance(images, list):
            return "Error: Please upload at least one image to proceed.", "Pending", "", "", 0

        milestones = []
        for image in images:
            # Ensure each image is a valid file path
            if not isinstance(image, str) or not os.path.isfile(image):
                return "Error: Invalid image file.", "Failure", "", "", 0

            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

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

            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:
                return f"Error: Failed to upload image to Salesforce - {str(e)}", "Failure", "", "", 0

            milestone = detect_milestone_from_image(saved_image_path)
            milestones.append(milestone)

        final_milestone = max(set(milestones), key=milestones.count) if milestones else "Planning"
        milestone_completion_map = {
            "Planning": 10,
            "Foundation": 30,
            "Walls Erected": 50,
            "Completed": 100,
        }
        percent_complete = milestone_completion_map.get(final_milestone, 0)

        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 f"<h3 style='color:green;'>Milestone: {final_milestone}</h3>", "Success", final_milestone, f"{percent_complete}%"

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

# Gradio UI with gallery input
with gr.Blocks(css="") as demo:
    gr.Markdown("<h1 class='title'>Construction Progress Analyzer</h1>")
    with gr.Row():
        image_input = gr.Gallery(
            label="Upload Construction Site Photos (JPG/PNG, ≤ 20MB)",
            type="filepath",
            interactive=True,
            show_download_button=False,
            show_share_button=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")
    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)