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
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib import colors

# 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)

# Inference function to generate captions dynamically based on image content
def generate_captions_from_image(image):
    if image.mode != "RGB":
        image = image.convert("RGB")
    
    # 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, 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
        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))
        
        # Build the PDF
        doc.build(flowables)
        return f"PDF saved successfully as {filename}"
    except Exception as e:
        return f"Error saving PDF: {str(e)}"

# Function to generate the daily progress report (DPR) text and save as PDF
def generate_dpr(files):
    dpr_text = []
    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 each uploaded file (image)
    for file in files:
        # Open the image from file path
        image = Image.open(file.name)  # Using file.name for filepath
        
        if image.mode != "RGB":
            image = image.convert("RGB")
        
        # Dynamically generate a caption based on the image
        caption = generate_captions_from_image(image)
        
        # Generate DPR section for this image with dynamic caption
        dpr_section = f"\nImage: {file.name}\nDescription: {caption}\n"
        dpr_text.append(dpr_section)
    
    # Combine DPR text
    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 = save_dpr_to_pdf(dpr_output, pdf_filename)
    
    # Return the DPR text and PDF save status
    return f"{dpr_output}\n\n{pdf_result}"

# Gradio interface for uploading multiple files and displaying the text-based DPR
iface = gr.Interface(
    fn=generate_dpr,
    inputs=gr.Files(type="filepath", label="Upload Site Photos"),  # Handle batch upload of images
    outputs="text",  # Display the DPR as text in the output section
    title="Daily Progress Report Generator",
    description="Upload up to 10 site photos. The AI model will dynamically detect construction activities, materials, and progress and generate a text-based Daily Progress Report (DPR). The DPR will also be saved as a PDF.",
    allow_flagging="never"  # Optional: Disable flagging
)

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