import os
import json
import base64
from google import genai
from google.genai import types

def ocr_scan_to_obsidian(image_path: str, obsidian_vault_path: str):
    """Uses Gemini to OCR a scanned notebook page and save it to Obsidian."""
    print(f"Starting OCR on {image_path}...")
    
    # Initialize the client
    client = genai.Client()
    
    # Read the image
    with open(image_path, "rb") as image_file:
        image_bytes = image_file.read()
        
    print("Uploading image to Gemini...")
    
    prompt = """
    You are an expert transcription assistant. 
    Carefully read the handwritten or printed notes in this scanned image.
    Convert the contents to high-quality Markdown.
    Preserve headings, bullet points, and any structural elements.
    If there are diagrams, describe them briefly.
    Output ONLY the markdown text.
    """
    
    response = client.models.generate_content(
        model='gemini-2.5-pro',
        contents=[
            prompt,
            types.Part.from_bytes(
                data=image_bytes,
                mime_type='image/png',
            ),
        ]
    )
    
    markdown_content = response.text.strip()
    
    # Create the obsidian note
    note_title = "Scanned Notebook Page.md"
    note_path = os.path.join(obsidian_vault_path, note_title)
    
    print(f"Saving to {note_path}...")
    
    # Add metadata frontmatter
    final_content = f"""---
source: scan
lineage_image: {image_path}
tags: [scanned, notebook]
---

{markdown_content}
"""

    # Ensure vault exists
    os.makedirs(obsidian_vault_path, exist_ok=True)
    
    with open(note_path, "w") as f:
        f.write(final_content)
        
    print("Successfully ingested scanned notebook to Obsidian.")

if __name__ == "__main__":
    scan_image = "/tmp/scan.png"
    vault = "/home/fcunha/desasossego"
    ocr_scan_to_obsidian(scan_image, vault)