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

# Setup logging
logging.basicConfig(
    filename="construction_analyzer.log",
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S"
)

# Load environment variables
logging.info("Loading 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]):
    logging.error("Missing Salesforce credentials")
    raise ValueError("Missing Salesforce credentials. Set SF_USERNAME, SF_PASSWORD, and SF_SECURITY_TOKEN in environment variables.")
logging.info("Salesforce credentials validated successfully")

# Initialize Salesforce connection
try:
    logging.info("Attempting Salesforce connection")
    sf = Salesforce(
        username=SF_USERNAME,
        password=SF_PASSWORD,
        security_token=SF_SECURITY_TOKEN,
        domain='login'
    )
    logging.info("Salesforce connection established successfully")
except Exception as e:
    logging.error(f"Salesforce connection failed: {str(e)}")
    raise

# Milestone definitions with completed and pending tasks
milestone_data = {
    "Excavation and Foundation": {
        "percentage": 10,
        "completed": ["Site clearing", "Excavation", "Foundation pouring"],
        "pending": ["Structural framework", "Roofing", "Exterior work", "Interior work", "Final inspection"]
    },
    "Structural Framework": {
        "percentage": 40,
        "completed": ["Site clearing", "Excavation", "Foundation pouring", "Structural columns and beams"],
        "pending": ["Roofing", "Exterior work", "Interior work", "Final inspection"]
    },
    "Roofing": {
        "percentage": 70,
        "completed": ["Site clearing", "Excavation", "Foundation pouring", "Structural columns and beams", "Roof installation"],
        "pending": ["Exterior work", "Interior work", "Final inspection"]
    },
    "Exterior Work": {
        "percentage": 85,
        "completed": ["Site clearing", "Excavation", "Foundation pouring", "Structural columns and beams", "Roof installation", "Exterior walls", "Windows and doors"],
        "pending": ["Interior work", "Final inspection"]
    },
    "Interior Work": {
        "percentage": 95,
        "completed": ["Site clearing", "Excavation", "Foundation pouring", "Structural columns and beams", "Roof installation", "Exterior walls", "Windows and doors", "Interior plumbing", "Electrical work", "Drywall and painting"],
        "pending": ["Final inspection"]
    },
    "Final Completion": {
        "percentage": 100,
        "completed": ["Site clearing", "Excavation", "Foundation pouring", "Structural columns and beams", "Roof installation", "Exterior walls", "Windows and doors", "Interior plumbing", "Electrical work", "Drywall and painting", "Final inspection"],
        "pending": []
    }
}

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

# Enhanced mock AI model simulating Grok-like analysis
def mock_ai_model(image):
    logging.info("Analyzing image for construction progress")
    img = image.convert("RGB")
    max_size = 1024
    img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)

    # Enhance contrast and brightness for feature detection
    enhancer = ImageEnhance.Contrast(img)
    img_enhanced = enhancer.enhance(2.0)
    enhancer = ImageEnhance.Brightness(img_enhanced)
    img_enhanced = enhancer.enhance(1.2)

    # Analyze image features
    img_data = list(img_enhanced.getdata())
    total_pixels = len(img_data)
    brightness_avg = sum(sum(pixel) / 3 for pixel in img_data) / total_pixels
    color_variation = max(max(pixel) - min(pixel) for pixel in img_data)

    # Edge detection
    edge_count = 0
    width, height = img_enhanced.size
    for x in range(width - 1):
        for y in range(height - 1):
            r, g, b = img_enhanced.getpixel((x, y))
            r_next, g_next, b_next = img_enhanced.getpixel((x + 1, y + 1))
            if abs(r - r_next) > 50 or abs(g - g_next) > 50 or abs(b - b_next) > 50:
                edge_count += 1
    edge_ratio = edge_count / (width * height)

    # Simulate Grok-like reasoning for milestone detection
    if brightness_avg > 220 and color_variation < 15 and edge_ratio < 0.05:
        milestone = "Final Completion"
        confidence = 0.95
    elif brightness_avg > 180 and edge_ratio < 0.1:
        milestone = "Interior Work"
        confidence = 0.90
    elif brightness_avg > 150 and edge_ratio < 0.2:
        milestone = "Exterior Work"
        confidence = 0.88
    elif brightness_avg > 120 and edge_ratio < 0.3:
        milestone = "Roofing"
        confidence = 0.85
    elif brightness_avg > 90 and edge_ratio < 0.4:
        milestone = "Structural Framework"
        confidence = 0.82
    else:
        milestone = "Excavation and Foundation"
        confidence = 0.80

    completed_tasks = milestone_data[milestone]["completed"]
    pending_tasks = milestone_data[milestone]["pending"]
    percentage = milestone_data[milestone]["percentage"]
    
    logging.info(f"Image analyzed: Milestone={milestone}, Percentage={percentage}, Confidence={confidence}")
    return milestone, percentage, confidence, completed_tasks, pending_tasks

# Process images and upload to Salesforce
def process_image(images, project_name, project_type):
    logging.info(f"Processing {len(images)} images for project {project_name}")
    try:
        if not images:
            logging.warning("No images uploaded")
            return "<p style='color: red;'>Error: Please upload at least one image.</p>", "Pending", "", "", 0

        if not project_name:
            logging.warning("Project name missing")
            return "<p style='color: red;'>Error: Project Name is required.</p>", "Pending", "", "", 0

        results = []
        image_urls = []
        all_percentages = []
        all_milestones = set()
        all_completed_tasks = set()
        all_pending_tasks = set()
        dominant_milestone = None
        dominant_image_url = None
        max_confidence = 0

        for idx, image_path in enumerate(images):
            try:
                img = Image.open(image_path)
                milestone, percent_complete, confidence, completed_tasks, pending_tasks = mock_ai_model(img)

                # Save image locally
                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}_{idx}_{os.path.basename(image_path)}"
                saved_image_path = os.path.join(upload_dir, image_filename)
                shutil.copy(image_path, saved_image_path)

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

                content_version = {
                    'Title': image_filename,
                    'PathOnClient': saved_image_path,
                    'VersionData': image_data
                }

                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}"
                    image_urls.append(file_url)
                    logging.info(f"Image {idx+1} uploaded to Salesforce: {file_url}")
                except Exception as e:
                    logging.error(f"Image {idx+1} upload failed: {str(e)}")
                    results.append(f"Image {idx+1}: Failed to upload to Salesforce - {str(e)}")
                    continue

                if percent_complete > (all_percentages[0] if all_percentages else -1):
                    dominant_milestone = milestone
                    dominant_image_url = file_url

                all_percentages.append(percent_complete)
                all_milestones.add(milestone)
                all_completed_tasks.update(completed_tasks)
                all_pending_tasks.update(pending_tasks)
                if confidence > max_confidence:
                    max_confidence = confidence

                results.append(
                    f"Image {idx+1}: {milestone} - {percent_complete}% completion (Confidence: {confidence})<br>"
                    f"<strong>Completed:</strong> {', '.join(completed_tasks)}<br>"
                    f"<strong>Pending:</strong> {', '.join(pending_tasks) if pending_tasks else 'None'}"
                )

            except Exception as e:
                logging.error(f"Image {idx+1} processing failed: {str(e)}")
                results.append(f"Image {idx+1}: Error processing image - {str(e)}")
                continue

        if not results:
            logging.warning("No images processed successfully")
            return "<p style='color: red;'>Error: No images were successfully processed.</p>", "Failure", "", "", 0

        total_percent_complete = round(sum(all_percentages) / len(all_percentages), 2)
        all_milestones_str = ", ".join(all_milestones)
        all_completed_tasks_str = ", ".join(sorted(all_completed_tasks))
        all_pending_tasks_str = ", ".join(sorted(all_pending_tasks)) if all_pending_tasks else "None"

        # Create Salesforce record
        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,
            "Project_Type__c": project_type,
            "Completion_Percentage__c": total_percent_complete,
            "Current_Milestone__c": all_milestones_str,
            "Last_Updated_On__c": local_time,
            "Upload_Status__c": "Success",
            "Comments__c": (
                f"Project {project_name} at {total_percent_complete}% completion. "
                f"Completed tasks: {all_completed_tasks_str}. "
                f"Pending tasks: {all_pending_tasks_str}."
            ),
            "Last_Updated_Image__c": dominant_image_url or ""
        }

        try:
            sf.Construction__c.create(record)
            logging.info(f"Salesforce record created for project {project_name}")
        except Exception as e:
            logging.error(f"Failed to create Salesforce record: {str(e)}")
            return f"<p style='color: red;'>Error: Failed to update Salesforce - {str(e)}</p>", "Failure", "", "", 0

        # Format output
        output_html = "<div class='output'>"
        output_html += "<h3>Processing Results:</h3><ul>"
        for result in results:
            if 'Error' in result or 'Failed' in result:
                output_html += f"<li class='error'>{result}</li>"
            else:
                output_html += f"<li class='success'>{result}</li>"
        output_html += "</ul>"
        output_html += f"<h3>Project Summary:</h3>"
        output_html += f"<p><strong>Project:</strong> {project_name} ({project_type})</p>"
        output_html += f"<p><strong>Total Completion:</strong> {total_percent_complete}%</p>"
        output_html += f"<p><strong>Milestones Detected:</strong> {all_milestones_str}</p>"
        output_html += f"<p><strong>Completed Tasks:</strong> {all_completed_tasks_str}</p>"
        output_html += f"<p><strong>Pending Tasks:</strong> {all_pending_tasks_str}</p>"
        output_html += f"<p><strong>Max Confidence:</strong> {max_confidence}</p>"
        output_html += f"<p><strong>Note:</strong> Only the image with the highest completion percentage is stored in Salesforce.</p>"
        output_html += "</div>"

        return output_html, "Success", "", f"Max Confidence: {max_confidence}", f"{total_percent_complete}%"

    except Exception as e:
        logging.error(f"Processing failed: {str(e)}")
        return f"<p style='color: red;'>Error: {str(e)}</p>", "Failure", "", "", "0%"

# Gradio UI
with gr.Blocks(css="""
    .gradio-container {
        background-color: #f9f9f9;
        font-family: 'Roboto', sans-serif;
        padding: 20px;
    }
    .title {
        color: #34495e;
        font-size: 30px;
        text-align: center;
        font-weight: bold;
        margin-bottom: 25px;
        text-transform: uppercase;
    }
    .gradio-row {
        text-align: center;
        max-width: 800px;
        margin: 0 auto;
    }
    .output {
        text-align: left;
        margin-top: 20px;
        padding: 30px;
        background-color: #ffffff;
        border-radius: 15px;
        box-shadow: 0 8px 16px rgba(0, 0, 0, 0.1);
        max-width: 800px;
        margin-left: auto;
        margin-right: auto;
    }
    .output h3 {
        color: #2c3e50;
        font-size: 22px;
        font-weight: bold;
        margin-bottom: 20px;
    }
    .output ul {
        list-style-type: none;
        padding: 0;
    }
    .output li {
        padding: 14px;
        margin-bottom: 18px;
        border-radius: 10px;
        font-size: 16px;
        transition: background-color 0.3s ease;
    }
    .output li.success {
        background-color: #27ae60;
        color: white;
    }
    .output li.error {
        background-color: #e74c3c;
        color: white;
    }
    .button {
        display: block;
        margin: 25px auto;
        background-color: #3498db;
        color: white;
        border: none;
        padding: 15px 30px;
        border-radius: 10px;
        cursor: pointer;
        font-size: 18px;
        transition: background-color 0.3s ease;
    }
    .button:hover {
        background-color: #2980b9;
    }
    .input {
        text-align: center;
        max-width: 800px;
        margin: 0 auto;
    }
""") as demo:
    gr.Markdown("<h1 class='title'>Construction Progress Analyzer</h1>")
    with gr.Row():
        project_type_input = gr.Dropdown(label="Project Type", choices=["House", "Apartment"], value="House")
        project_name_input = gr.Textbox(label="Project Name (Required)", placeholder="e.g. Project_12345")
    
    image_input = gr.File(
        file_count="multiple",
        file_types=["image"],
        label="Upload Construction Site Photos (JPG/PNG, ≤20MB each)"
    )
    submit_button = gr.Button("Process Images")
    output_html = gr.HTML(label="Result")
    
    submit_button.click(
        fn=process_image,
        inputs=[image_input, project_name_input, project_type_input],
        outputs=[output_html]
    )

logging.info("Launching Gradio interface")
demo.launch()