app.py

from fastapi import FastAPI, File, UploadFile
import requests

from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image
import gradio as gr
import torch
from datetime import datetime
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image as PDFImage
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib import colors
from simple_salesforce import Salesforce
import os
from dotenv import load_dotenv
import base64
import io
import concurrent.futures

# Load environment variables from .env file
load_dotenv()

app = FastAPI()


# Salesforce credentials
SF_USERNAME = os.getenv('SF_USERNAME')
SF_PASSWORD = os.getenv('SF_PASSWORD')
SF_SECURITY_TOKEN = os.getenv('SF_SECURITY_TOKEN')

# Initialize Salesforce connection
try:
    sf = Salesforce(username=SF_USERNAME, password=SF_PASSWORD, security_token=SF_SECURITY_TOKEN)
except Exception as e:
    sf = None
    print(f"Failed to connect to Salesforce: {str(e)}")

# Load BLIP model and processor
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
model.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)


# FastAPI endpoint to handle image upload and caption generation
@app.post("/predict/")
async def predict(image: UploadFile = File(...)):
    try:
        # Read the image from the request
        image_bytes = await image.read()
        image = Image.open(BytesIO(image_bytes))
        
        # Generate caption from the image
        caption = generate_captions_from_image(image)
        return {"caption": caption}
    except Exception as e:
        return {"error": str(e)}


# Inference function to generate captions dynamically based on image content
def generate_captions_from_image(image):
    if image.mode != "RGB":
        image = image.convert("RGB")
    
    # Resize image for faster processing
    image = image.resize((640, 640))
    
    # Preprocess the image and generate a caption
    inputs = processor(image, return_tensors="pt").to(device, torch.float16)
    output = model.generate(**inputs, max_new_tokens=50)
    caption = processor.decode(output[0], skip_special_tokens=True)
    
    return caption

# Function to save DPR text to a PDF file
def save_dpr_to_pdf(dpr_text, image_paths, captions, filename):
    try:
        # Create a PDF document
        doc = SimpleDocTemplate(filename, pagesize=letter)
        styles = getSampleStyleSheet()
        
        # Define custom styles
        title_style = ParagraphStyle(
            name='Title',
            fontSize=16,
            leading=20,
            alignment=1,  # Center
            spaceAfter=20,
            textColor=colors.black,
            fontName='Helvetica-Bold'
        )
        body_style = ParagraphStyle(
            name='Body',
            fontSize=12,
            leading=14,
            spaceAfter=10,
            textColor=colors.black,
            fontName='Helvetica'
        )
        
        # Build the PDF content
        flowables = []
        
        # Add title
        flowables.append(Paragraph("Daily Progress Report", title_style))
        
        # Split DPR text into lines and add as paragraphs (excluding descriptions for images)
        for line in dpr_text.split('\n'):
            # Replace problematic characters for PDF
            line = line.replace('\u2019', "'").replace('\u2018', "'")
            if line.strip():
                flowables.append(Paragraph(line, body_style))
            else:
                flowables.append(Spacer(1, 12))
        
        # Add images and captions in the correct order (no need to add description to dpr_text again)
        for img_path, caption in zip(image_paths, captions):
            try:
                # Add image first
                img = PDFImage(img_path, width=200, height=150)  # Adjust image size if needed
                flowables.append(img)
                # Add description below the image
                description = f"Description: {caption}"
                flowables.append(Paragraph(description, body_style))
                flowables.append(Spacer(1, 12))  # Add some space between images
            except Exception as e:
                flowables.append(Paragraph(f"Error loading image: {str(e)}", body_style))
        
        # Build the PDF
        doc.build(flowables)
        return f"PDF saved successfully as {filename}", filename
    except Exception as e:
        return f"Error saving PDF: {str(e)}", None

# Function to upload a file to Salesforce as ContentVersion
def upload_file_to_salesforce(file_path, filename, sf_connection, file_type):
    try:
        # Read file content and encode in base64
        with open(file_path, 'rb') as f:
            file_content = f.read()
        file_content_b64 = base64.b64encode(file_content).decode('utf-8')
        
        # Set description based on file type
        description = "Daily Progress Report PDF" if file_type == "pdf" else "Site Image"
        
        # Create ContentVersion
        content_version = sf_connection.ContentVersion.create({
            'Title': filename,
            'PathOnClient': filename,
            'VersionData': file_content_b64,
            'Description': description
        })
        
        # Get ContentDocumentId
        content_version_id = content_version['id']
        content_document = sf_connection.query(
            f"SELECT ContentDocumentId FROM ContentVersion WHERE Id = '{content_version_id}'"
        )
        content_document_id = content_document['records'][0]['ContentDocumentId']
        
        # Generate a valid Salesforce URL for the ContentDocument
        content_document_url = f"https://{sf_connection.sf_instance}.salesforce.com/{content_document_id}"

        # Ensure the link is valid
        return content_document_id, content_document_url, f"File {filename} uploaded successfully"
    except Exception as e:
        return None, None, f"Error uploading {filename} to Salesforce: {str(e)}"

# Function to generate the daily progress report (DPR), save as PDF, and upload to Salesforce
def generate_dpr(files):
    dpr_text = []
    captions = []
    image_paths = []
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # Add header to the DPR
    dpr_text.append(f"Daily Progress Report\nGenerated on: {current_time}\n")
    
    # Process images in parallel for faster performance
    with concurrent.futures.ThreadPoolExecutor() as executor:
        results = list(executor.map(lambda file: generate_captions_from_image(Image.open(file.name)), files))
    
    for i, file in enumerate(files):
        caption = results[i]
        captions.append(caption)
        
        # Generate DPR section for this image with dynamic caption
        dpr_section = f"\nImage: {file.name}\nDescription: {caption}\n"
        # Remove the description from the dpr_text section
        # No need to add it again as the image and caption will be inserted in the PDF
        dpr_text.append(dpr_section)
        
        # Save image path for embedding in the report
        image_paths.append(file.name)
    
    # Combine DPR text (no redundant description here)
    dpr_output = "\n".join(dpr_text)
    
    # Generate PDF filename with timestamp
    pdf_filename = f"DPR_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.pdf"
    
    # Save DPR text to PDF
    pdf_result, pdf_filepath = save_dpr_to_pdf(dpr_output, image_paths, captions, pdf_filename)
    
    # Salesforce upload
    salesforce_result = ""
    pdf_content_document_id = None
    pdf_url = None
    image_content_document_ids = []
    
    if sf and pdf_filepath:
        try:
            # Create Daily_Progress_Reports__c record
            report_description = "; ".join(captions)[:255]  # Concatenate captions, limit to 255 chars
            dpr_record = sf.Daily_Progress_Reports__c.create({
                'Detected_Activities__c': report_description  # Store in Detected_Activities__c field
            })
            dpr_record_id = dpr_record['id']
            salesforce_result += f"Created Daily_Progress_Reports__c record with ID: {dpr_record_id}\n"
            
            # Upload PDF to Salesforce
            pdf_content_document_id, pdf_url, pdf_upload_result = upload_file_to_salesforce(
                pdf_filepath, pdf_filename, sf, "pdf"
            )
            salesforce_result += pdf_upload_result + "\n"
            
            # Link PDF to DPR record
            if pdf_content_document_id:
                sf.ContentDocumentLink.create({
                    'ContentDocumentId': pdf_content_document_id,
                    'LinkedEntityId': dpr_record_id,
                    'ShareType': 'V'
                })
            
            # Update the DPR record with the PDF URL
            if pdf_url:
                sf.Daily_Progress_Reports__c.update(dpr_record_id, {
                    'PDF_URL__c': pdf_url  # Storing the PDF URL correctly
                })
                salesforce_result += f"Updated PDF URL for record ID {dpr_record_id}\n"
            
            # Upload images to Salesforce and create Site_Images__c records
            for file in files:
                image_filename = os.path.basename(file.name)
                image_content_document_id, image_upload_result = upload_file_to_salesforce(
                    file.name, image_filename, sf, "image"
                )
                if image_content_document_id:
                    image_content_document_ids.append(image_content_document_id)
                
                # Create Site_Images__c record and link to DPR
                site_image_record = sf.Site_Images__c.create({
                    'Image__c': image_content_document_id,
                    'Related_Report__c': dpr_record_id  # Link image to DPR record
                })
                salesforce_result += image_upload_result + "\n"
                
                # Link image to DPR record
                if image_content_document_id:
                    sf.ContentDocumentLink.create({
                        'ContentDocumentId': image_content_document_id,
                        'LinkedEntityId': dpr_record_id,
                        'ShareType': 'V'
                    })
            
        except Exception as e:
            salesforce_result += f"Error interacting with Salesforce: {str(e)}\n"
    else:
        salesforce_result = "Salesforce connection not available or PDF generation failed.\n"
    
    # Return DPR text, PDF file, and Salesforce upload status
    return (
        dpr_output + f"\n\n{pdf_result}\n\nSalesforce Upload Status:\n{salesforce_result}",
        pdf_filepath
    )

# Gradio interface for uploading multiple files, displaying DPR, and downloading PDF
iface = gr.Interface(
    fn=generate_dpr,
    inputs=gr.Files(type="filepath", label="Upload Site Photos"),
    outputs=[
        gr.Textbox(label="Daily Progress Report"),
        gr.File(label="Download PDF")
    ],
    title="Daily Progress Report Generator",
    description="Upload up to 10 site photos. The AI model will generate a text-based Daily Progress Report (DPR), save it as a PDF, and upload the PDF and images to Salesforce under Daily_Progress_Reports__c in the Files related list. Download the PDF locally if needed.",
    allow_flagging="never"
)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)