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

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

# Deterministic AI prediction with fixed confidence and percent
def mock_ai_model(image):
    img = image.convert("RGB")
    max_size = 1024
    img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)

    img_bytes = img.tobytes()
    img_hash = int(hashlib.sha256(img_bytes).hexdigest(), 16)

    milestone_index = img_hash % len(VALID_MILESTONES)
    milestone = VALID_MILESTONES[milestone_index]

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

    confidence_raw = 0.85 + ((img_hash % 1000) / 1000) * (0.95 - 0.85)
    confidence_score = round(confidence_raw, 2)

    return milestone, completion_percent, confidence_score

# Image processing and Salesforce upload
def process_image(image, project_name):
    try:
        if image is None:
            return "Error: Please upload an image to proceed.", "Pending", "", "", 0

        img = Image.open(image)
        image_size_mb = os.path.getsize(image) / (1024 * 1024)
        if image_size_mb > 20:
            return "Error: Image size exceeds 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
        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)

        # Corrected public URL logic
        if os.getenv("GRADIO_SERVER_NAME"):
            public_url_base = f"https://{os.getenv('GRADIO_SERVER_NAME')}/file"
        else:
            public_url_base = "http://localhost:7860/file"

        image_url = f"{public_url_base}/{upload_dir}/{image_filename}"

        milestone, percent_complete, confidence_score = mock_ai_model(img)

        record = {
            "Name__c": project_name,
            "Current_Milestone__c": milestone,
            "Completion_Percentage__c": percent_complete,
            "Last_Updated_On__c": datetime.now().isoformat(),
            "Upload_Status__c": "Success",
            "Comments__c": f"AI Prediction: {milestone} with {confidence_score*100}% confidence",
            "Last_Updated_Image__c": image_url
        }

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

        return (
            f"Success: Milestone: {milestone}, Completion: {percent_complete}%",
            "Success",
            milestone,
            f"Confidence Score: {confidence_score}",
            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;} .title {color: #2c3e50; font-size: 24px; text-align: center;}") as demo:
    gr.Markdown("<h1 class='title'>Construction Milestone Detector</h1>")
    with gr.Row():
        image_input = gr.Image(type="filepath", label="Upload Construction Site Photo (JPG/PNG, ≤ 20MB)")
        project_name_input = gr.Textbox(label="Project Name (Required)", placeholder="e.g. Project_12345")

    submit_button = gr.Button("Process Image")
    output_text = gr.Textbox(label="Result")
    upload_status = gr.Textbox(label="Upload Status")
    milestone = gr.Textbox(label="Detected Milestone")
    confidence = gr.Textbox(label="Confidence Score")
    progress = gr.Slider(0, 100, label="Completion Percentage", interactive=False, value=0)

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

demo.launch(share=True)