import gradio as gr
from analyzer import analyze_cells
from simple_salesforce import Salesforce
import io
import base64
import tempfile
from PIL import Image

# Salesforce credentials and connection
sf_username = 'scope.analyzer@smartlabtech.com'
sf_password = 'Scope@1020'
sf_security_token = 'IPbKc5V7VwT96RMWAVgy3Y7ol'
sf_instance_url = 'https://sathkruthatechsolutionsp-2f-dev-ed.develop.lightning.force.com'

# Create Salesforce connection
sf = Salesforce(username=sf_username, password=sf_password, security_token=sf_security_token)

def upload_image_to_salesforce(image_data, image_name, record_id=None):
    """Upload the image to Salesforce as a ContentVersion and return a public URL."""
    try:
        encoded_image_data = base64.b64encode(image_data).decode('utf-8')
        content_version_data = {
            "Title": image_name,
            "PathOnClient": image_name,
            "VersionData": encoded_image_data,
        }
        if record_id:
            content_version_data["FirstPublishLocationId"] = record_id
        content_version = sf.ContentVersion.create(content_version_data)
        content_version_id = content_version["id"]
        image_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version_id}"
        return image_url
    except Exception as e:
        print(f"Error uploading image to Salesforce: {e}")
        return None

def create_salesforce_record(cell_count, original_image_url, annotated_image_url, anomaly_type):
    try:
        record = {
            'Cell_Count__c': cell_count,
            'Image_Original__c': original_image_url,
            'Annotated_Image__c': annotated_image_url,
            'Anomaly_Type__c': anomaly_type  # Add anomaly type
        }
        response = sf.Cell_Analysis__c.create(record)
        print(f"Record created in Salesforce with ID: {response['id']}")
        return f"Cell Count: {cell_count}"  # Only return the cell count message (without "Anomaly: ")
    except Exception as e:
        print(f"Error creating record in Salesforce: {e}")
        return f"Error creating record in Salesforce: {e}"

def analyze_image(uploaded_file):
    try:
        count, annotated_image, original_image, anomaly_type = analyze_cells(uploaded_file)
        if count is None:
            return None, "Anomaly detected: The uploaded image is too blurry or unclear.", "", ""

        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as annotated_image_file:
            annotated_image.save(annotated_image_file, format="JPEG")
            annotated_image_path = annotated_image_file.name

        original_image_url = upload_image_to_salesforce(original_image, "Original_Image.jpg")
        annotated_image_url = upload_image_to_salesforce(open(annotated_image_path, "rb").read(), "Annotated_Image.jpg")
        response_message = create_salesforce_record(count, original_image_url, annotated_image_url, anomaly_type)

        return annotated_image, response_message, anomaly_type  # Return anomaly type separately
    except Exception as e:
        print(f"Error analyzing image: {e}")
        return None, f"Error analyzing image: {e}", "", ""

# Create the Gradio Interface
with gr.Blocks() as app:
    gr.HTML("<h1 style='text-align:center; color: #4A90E2;'>Microscope Image Analyzer</h1>")
    
    with gr.Row():
        with gr.Column():
            uploaded_file = gr.File(file_types=[".png", ".jpg", ".jpeg"], type="binary", label="Upload Image", interactive=True)
            with gr.Row():
                clear_btn = gr.Button("Clear", variant="secondary")
                submit_btn = gr.Button("Submit", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="Annotated Image", type="pil")
            output_text = gr.Textbox(label="Cell Count", interactive=False)
            output_anomaly = gr.Textbox(label="Anomaly Type", interactive=False)  # Display anomaly type separately

    submit_btn.click(
        fn=analyze_image,
        inputs=uploaded_file,
        outputs=[output_image, output_text, output_anomaly]  # Anomaly type is now separate
    )

    clear_btn.click(
        fn=lambda: [None, "", ""],
        outputs=[output_image, output_text, output_anomaly]  # Clear both outputs
    )

app.launch(share=False)