import gradio as gr
from PIL import Image, ImageDraw
import torch
from torchvision import models, transforms
from simple_salesforce import Salesforce
import base64
from io import BytesIO
import logging
from datetime import datetime

# Setup logging
logging.basicConfig(level=logging.INFO)

# Salesforce Credentials (ensure these are set up in your environment variables for security)
SALESFORCE_USERNAME = "drone@sathkrutha.com"
SALESFORCE_PASSWORD = "Komal1303@"
SALESFORCE_SECURITY_TOKEN = "53AWRskW9EjWUsSL5LU6nFTy3"
SALESFORCE_INSTANCE_URL = "https://sathikrutha-a-dev-ed.my.salesforce.com"

# Replace with a valid Site__c record ID from your Salesforce org
SITE_RECORD_ID = "a003000000xxxxx"  # TODO: Update with actual ID from Site__c

# Connect to Salesforce
try:
    sf = Salesforce(
        username=SALESFORCE_USERNAME,
        password=SALESFORCE_PASSWORD,
        security_token=SALESFORCE_SECURITY_TOKEN,
        instance_url=SALESFORCE_INSTANCE_URL
    )
    logging.info("Salesforce connection established.")
except Exception as e:
    logging.error(f"Failed to connect to Salesforce: {str(e)}")
    raise Exception(f"Failed to connect to Salesforce: {str(e)}")

# Load Model (updated to use correct weights for the current PyTorch version)
model = models.detection.fasterrcnn_resnet50_fpn(weights="FasterRCNN_ResNet50_FPN_Weights.COCO_V1")
model.eval()

# Define labels (COCO labels; fine-tune for structural defects)
COCO_INSTANCE_CATEGORY_NAMES = [
    '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter',
    'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra',
    'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
    'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
    'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon',
    'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet',
    'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
    'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
    'hair drier', 'toothbrush'
]

# Image transformations
transform = transforms.Compose([
    transforms.ToTensor(),
])

# Map model severity to Salesforce picklist values
def get_severity(score):
    if score >= 0.9:
        return "Critical"
    elif score >= 0.7:
        return "Moderate"
    else:
        return "Minor"

# Temporary mapping for COCO labels to structural defects
COCO_TO_DEFECT_MAPPING = {
    'car': 'Crack',
    'person': 'Rust',
    'bicycle': 'Deformation',
    'truck': 'Corrosion',
    'boat': 'Spalling',
}

def map_defect_type(coco_label):
    return COCO_TO_DEFECT_MAPPING.get(coco_label, "Crack")

# Function to upload image to Salesforce as ContentVersion
def upload_image_to_salesforce(image, filename="detected_image.jpg", record_id=None):
    try:
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        img_data = base64.b64encode(buffered.getvalue()).decode("utf-8")
        
        content_version = sf.ContentVersion.create({
            "Title": filename,
            "PathOnClient": filename,
            "VersionData": img_data,
            "FirstPublishLocationId": record_id if record_id else None
        })
        logging.info(f"Image uploaded to Salesforce with ContentVersion ID: {content_version['id']}")
        return content_version["id"]
    except Exception as e:
        logging.error(f"Failed to upload image to Salesforce: {str(e)}")
        raise Exception(f"Failed to upload image to Salesforce: {str(e)}")

# Detect defects and integrate with Salesforce
def detect_defects(image):
    if not image:
        return None, {"error": "No image provided"}

    try:
        # Perform detection
        image_tensor = transform(image).unsqueeze(0)
        with torch.no_grad():
            predictions = model(image_tensor)

        result_image = image.copy()
        draw = ImageDraw.Draw(result_image)
        output = []

        for i in range(len(predictions[0]['boxes'])):
            score = predictions[0]['scores'][i].item()
            if score < 0.3:  # Lowered threshold to detect more objects
                continue

            box = predictions[0]['boxes'][i].tolist()
            label_idx = predictions[0]['labels'][i].item()
            coco_label = COCO_INSTANCE_CATEGORY_NAMES[label_idx]
            defect_type = map_defect_type(coco_label)
            severity = get_severity(score)

            output.append({
                "type": defect_type,
                "confidence": round(score, 2),
                "severity": severity,
                "coco_label": coco_label
            })

            draw.rectangle(box, outline="red", width=3)
            draw.text((box[0], box[1]), f"{defect_type}: {severity}", fill="red")

        # Create Salesforce record if detections exist
        if output:
            try:
                current_date = datetime.now().strftime("%Y-%m-%d")
                inspection_name = f"Inspection-{current_date}-{len(output):03d}"

                # Creating the Salesforce record with updated fields
                inspection_record = sf.Drone_Structure_Inspection__c.create({
                    "Inspection_Date__c": current_date,
                    "Fault_Type__c": output[0]["type"],  # Mapping defect type
                    "Severity__c": output[0]["severity"],  # Mapping severity
                    "Fault_Summary__c": str(output),  # Summarizing the defects
                    "Status__c": "New",  # Default status
                    "Annotated_Image_URL__c": "",  # Placeholder for image URL
                    "Report_PDF__c": ""  # Placeholder for report PDF URL
                })

                record_id = inspection_record.get("id")
                content_version_id = upload_image_to_salesforce(
                    result_image,
                    filename=f"detected_defect_{record_id}.jpg",
                    record_id=record_id
                )

                if content_version_id:
                    sf.Drone_Structure_Inspection__c.update(record_id, {
                        "Annotated_Image_URL__c": f"/sfc/servlet.shepherd/version/download/{content_version_id}"
                    })

                output.append({"salesforce_record_id": record_id})
            except Exception as e:
                output.append({"error": f"Failed to create Salesforce record: {str(e)}"})

        return result_image, output

    except Exception as e:
        logging.error(f"Processing failed: {str(e)}")
        return None, {"error": f"Processing failed: {str(e)}"}

# Gradio Interface
demo = gr.Interface(
    fn=detect_defects,
    inputs=gr.Image(type="pil", label="Upload Drone Image"),
    outputs=[
        gr.Image(label="Detection Result"),
        gr.Textbox(label="Detected Faults with Severity")
    ],
    title="Structural Defect Detection with Salesforce Integration",
    description="Detects objects using Faster R-CNN and stores results in Salesforce. Fine-tune the model for structural defects like cracks, rust, and spalling."
)

if __name__ == "__main__":
    demo.launch()