program.py

import os
import json

from pathlib import Path
from dotenv import load_dotenv
from PIL import Image, ExifTags

# google gemini
from google import genai

# pdf
from reportlab.lib.pagesizes import A4
from reportlab.lib import colors
from reportlab.lib.units import cm
from reportlab.platypus import SimpleDocTemplate, Table, TableStyle, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.enums import TA_CENTER

# excel
from openpyxl import Workbook
from openpyxl.styles import Font, PatternFill, Alignment, Border, Side, PatternFill
from openpyxl.utils import get_column_letter

# word
from docx import Document
from docx.shared import Pt, RGBColor, Cm, Inches
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.enum.table import WD_ALIGN_VERTICAL
from docx.oxml.ns import qn
from docx.oxml import OxmlElement


def simulate_gemini_response():
    """
    Mock response for development purposes.
    Avoids unnecessary API calls during testing.
    Replace ocr_handwriting() with this function to run offline.
    """    
    mock_json_string = '[["Item", "check", "Notes"], ["Water quality", "Y", "OK"], ["Air Filters", "N", "Replace"], ["Temperature", "Y", "OK"], ["Humidity", "Y", "OK"], ["Waste bins", "N", "Full"]]'
    
    return json.loads(mock_json_string)


load_dotenv()
API_KEY = os.getenv("GEMINI_API_KEY")
# MODEL = "gemini-2.5-flash"
MODEL = "gemini-2.5-flash-lite"

BASE_DIR = Path(__file__).parent
INPUT_DIR = BASE_DIR / "input"
OUTPUT_DIR = BASE_DIR / "output"

DEBUG = True  # Set to False to run the full pipeline with a real image and API call


def gen_pdf(data, file_name):

    DOC_MARGIN = 2*cm
    doc = SimpleDocTemplate(
        str(file_name),
        pagesize=A4,
        leftMargin=DOC_MARGIN,
        rightMargin=DOC_MARGIN,
        topMargin=DOC_MARGIN,
        bottomMargin=DOC_MARGIN
    )
    doc.title = "OCR Report"
    doc.author = "Handwritten OCR Pipeline"

    # Cores (consistentes com o Excel)
    HEADER_BG   = colors.HexColor("#2D4A8A")
    ROW_ALT     = colors.HexColor("#EEF2FA")
    ROW_WHITE   = colors.white
    BORDER      = colors.HexColor("#BFCCE8")

    styles = getSampleStyleSheet()

    title_style = ParagraphStyle(
        "title",
        parent=styles["Title"],
        fontName="Helvetica-Bold",
        fontSize=16,
        textColor=HEADER_BG,
        spaceAfter=0.5*cm,
        alignment=TA_CENTER
    )

    subtitle_style = ParagraphStyle(
        "subtitle",
        parent=styles["Normal"],
        fontName="Helvetica",
        fontSize=9,
        textColor=colors.grey,
        spaceAfter=1*cm,
        alignment=TA_CENTER
    )

    story = []

    story.append(Paragraph("OCR Report", title_style))
    story.append(Paragraph("Generated automatically from handwritten image", subtitle_style))

    headers = data[0]
    rows = data[1:]

    table_style_commands = [
        # Header
        ("BACKGROUND",      (0, 0), (-1, 0),  HEADER_BG),
        ("TEXTCOLOR",       (0, 0), (-1, 0),  colors.white),
        ("FONTNAME",        (0, 0), (-1, 0),  "Helvetica-Bold"),
        ("FONTSIZE",        (0, 0), (-1, 0),  11),
        ("ROWBACKGROUND",   (0, 0), (-1, 0), HEADER_BG),

        # Data
        ("FONTNAME",        (0, 1), (-1, -1), "Helvetica"),
        ("FONTSIZE",        (0, 1), (-1, -1), 10),
        ("TEXTCOLOR",       (0, 1), (-1, -1), colors.HexColor("#222222")),

        # Alignment
        ("ALIGN",           (0, 0), (-1, -1), "CENTER"),
        ("VALIGN",          (0, 0), (-1, -1), "MIDDLE"),

        # Padding
        ("TOPPADDING",      (0, 0), (-1, -1), 8),
        ("BOTTOMPADDING",   (0, 0), (-1, -1), 8),

        # Borders
        ("GRID",            (0, 0), (-1, -1), 0.5, BORDER),
        ("BOX",             (0, 0), (-1, -1), 1,   HEADER_BG),

        # ALTERNATE ROW BACKGROUND COLORS
        ("ROWBACKGROUNDS",  (0, 1), (-1, -1), [ROW_ALT, ROW_WHITE]),
    ]


    # Cols width
    page_width = A4[0] - 2 * DOC_MARGIN  # left + right margins
    col_width = page_width / len(headers)
    col_widths = [col_width] * len(headers)

    # Create table
    table = Table([headers] + rows, colWidths=col_widths, repeatRows=1)
    table.setStyle(TableStyle(table_style_commands))

    story.append(table)

    # footnote
    story.append(Spacer(1, 0.8*cm))
    story.append(Paragraph(
        "* Data extracted via OCR using Google Gemini API",
        ParagraphStyle("note", parent=styles["Normal"], fontSize=8, textColor=colors.grey)
    ))

    doc.build(story)

def gen_excel(data, file_name):
    wb = Workbook()
    ws = wb.active
    ws.title = "OCR Report"

    # Colors
    HEADER_BG = "2D4A8A"   # dark blue
    HEADER_FG = "FFFFFF"   # white
    ROW_ALT   = "EEF2FA"   # light blue
    ROW_WHITE = "FFFFFF"   # white
    BORDER_COLOR = "BFCCE8"

    thin = Side(style="thin", color=BORDER_COLOR)
    border = Border(left=thin, right=thin, top=thin, bottom=thin)

    headers = data[0]
    rows = data[1:]

    # Headers
    for col_idx, header in enumerate(headers, start=1):
        cell = ws.cell(row=1, column=col_idx, value=header)
        cell.font = Font(name="Arial", bold=True, color=HEADER_FG, size=11)
        cell.fill = PatternFill("solid", start_color=HEADER_BG)
        cell.alignment = Alignment(horizontal="center", vertical="center")
        cell.border = border

    ws.row_dimensions[1].height = 30

    # Data
    for row_idx, row in enumerate(rows, start=2):
        bg = ROW_ALT if row_idx % 2 == 0 else ROW_WHITE 
        for col_idx, value in enumerate(row, start=1):
            cell = ws.cell(row=row_idx, column=col_idx, value=value)
            cell.font = Font(name="Arial", size=11)
            cell.fill = PatternFill("solid", start_color=bg)
            cell.alignment = Alignment(horizontal="center", vertical="center")
            cell.border = border
        ws.row_dimensions[row_idx].height = 22

    # cols width
    for col_idx, _ in enumerate(headers, start=1): # for every column
        col_letter = get_column_letter(col_idx)
        
        max_content_length = 0
        total_rows = len(data)
        
        for row_idx in range(1, total_rows + 1): # for every row
            cell_value = ws.cell(row=row_idx, column=col_idx).value
            cell_text = str(cell_value) if cell_value is not None else ""
            cell_length = len(cell_text)
            
            if cell_length > max_content_length: # get max length
                max_content_length = cell_length 
        
        # Set width
        ws.column_dimensions[col_letter].width = max_content_length + 6

    # Freeze header row
    ws.freeze_panes = "A2"

    wb.save(file_name)

def gen_word(data, file_name):
    # COLORS
    HEADER_BG  = "2D4A8A"
    ROW_ALT    = "EEF2FA"
    ROW_WHITE  = "FFFFFF"

    def set_cell_background(cell, hex_color):
        """Cell background color."""
        tc = cell._tc
        tcPr = tc.get_or_add_tcPr()
        shd = OxmlElement("w:shd")
        shd.set(qn("w:fill"), hex_color)
        shd.set(qn("w:val"), "clear")
        tcPr.append(shd)    

    doc = Document()

    # Margin
    for section in doc.sections:
        section.top_margin    = Cm(2)
        section.bottom_margin = Cm(2)
        section.left_margin   = Cm(2)
        section.right_margin  = Cm(2)

    # Title
    title = doc.add_paragraph("OCR Report")
    title.alignment = WD_ALIGN_PARAGRAPH.CENTER
    title_run = title.runs[0]
    title_run.font.name      = "Arial"
    title_run.font.bold      = True
    title_run.font.size      = Pt(18)
    title_run.font.color.rgb = RGBColor(0x2D, 0x4A, 0x8A)

    # Sub Title
    subtitle = doc.add_paragraph("Generated automatically from handwritten image")
    subtitle.alignment = WD_ALIGN_PARAGRAPH.CENTER
    subtitle_run = subtitle.runs[0]
    subtitle_run.font.name      = "Arial"
    subtitle_run.font.size      = Pt(9)
    subtitle_run.font.color.rgb = RGBColor(0x88, 0x88, 0x88)

    doc.add_paragraph()  # Add empty paragraph for space

    headers = data[0]
    rows    = data[1:]

    # Create table
    num_cols = len(headers)
    num_rows = len(data)
    table = doc.add_table(rows=num_rows, cols=num_cols)
    table.style = "Table Grid"

    # Cols width
    page_width   = Inches(6.5)  
    col_width    = page_width / num_cols

    # Header row
    header_row = table.rows[0]
    header_row.height = Cm(1)

    for col_idx, header_text in enumerate(headers):
        cell = header_row.cells[col_idx]
        cell.width = col_width
        set_cell_background(cell, HEADER_BG)

        paragraph = cell.paragraphs[0]
        paragraph.alignment = WD_ALIGN_PARAGRAPH.CENTER
        cell.vertical_alignment = WD_ALIGN_VERTICAL.CENTER

        run = paragraph.add_run(header_text)
        run.font.name      = "Arial"
        run.font.bold      = True
        run.font.size      = Pt(11)
        run.font.color.rgb = RGBColor(0xFF, 0xFF, 0xFF)

    # data rows
    for row_idx, row_data in enumerate(rows, start=1):
        bg_color = ROW_WHITE if row_idx % 2 == 0 else ROW_ALT
        word_row = table.rows[row_idx]
        word_row.height = Cm(0.8)

        for col_idx, cell_value in enumerate(row_data):
            cell = word_row.cells[col_idx]
            cell.width = col_width
            set_cell_background(cell, bg_color)

            paragraph = cell.paragraphs[0]
            paragraph.alignment = WD_ALIGN_PARAGRAPH.CENTER
            cell.vertical_alignment = WD_ALIGN_VERTICAL.CENTER

            run = paragraph.add_run(str(cell_value))
            run.font.name = "Arial"
            run.font.size = Pt(10)

    # footnote
    doc.add_paragraph()
    note = doc.add_paragraph("* Data extracted via OCR using Google Gemini API")
    note_run = note.runs[0]
    note_run.font.name      = "Arial"
    note_run.font.size      = Pt(8)
    note_run.font.color.rgb = RGBColor(0x88, 0x88, 0x88)

    doc.save(str(file_name))



def fix_rotation(image_path):
    img = Image.open(image_path)
    
    try:
        exif = img._getexif()
        if exif:
            for tag, value in exif.items():
                if ExifTags.TAGS.get(tag) == 'Orientation':
                    if value == 3:
                        img = img.rotate(180, expand=True)
                    elif value == 6:
                        img = img.rotate(270, expand=True)
                    elif value == 8:
                        img = img.rotate(90, expand=True)
    except:
        pass
    
    return img

def ocr_handwriting(image_path):
    try:
        client = genai.Client(api_key=API_KEY, http_options={'api_version': 'v1'})
        # img = Image.open(image_path)
        img = fix_rotation(image_path)
        # img.save(OUTPUT_DIR / "debug_rotated.jpeg") # [DEBUG]
        
        # List available models
        print(f"Available models:\n{', \n'.join([f"{idx}) {x.name} # {x.display_name}" for idx, x in enumerate(client.models.list(), start=1)])}")
        print()
        print(f"Processing file with model {MODEL}: {image_path}...\n")
        
        query = """
        Read the table in the image.
        Return ONLY a raw JSON array of arrays. No markdown, no extra text.
        First inner array is the headers, remaining are data rows.
        Example: [["Col1", "Col2"], ["val1", "val2"], ["val3", "val4"]]
        """

        response = client.models.generate_content(
            model=MODEL,
            contents=[query, img]
        )

        result = json.loads(response.text)

        return result
        
    except Exception as e:
        return f"Error: {e}"


def main():
    filePath = INPUT_DIR / "handwrite.jpeg"
    
    if not filePath.exists():
        print(f"Error: File {filePath} not found.")
        return

    if DEBUG:
        result = simulate_gemini_response() # Mock data for development
    else:
        result = ocr_handwriting(str(filePath)) # Live API call

    # print("\n--- RESULT ---")
    # print(f"List: {result}")

    gen_pdf(result, OUTPUT_DIR / "report.pdf")
    gen_excel(result, OUTPUT_DIR / "report.xlsx")
    gen_word(result, OUTPUT_DIR / "report.docx")

if __name__ == "__main__":
    main()