import os
import io
from google.api_core.client_options import ClientOptions
from google.cloud import documentai
from dotenv import load_dotenv

load_dotenv()

PROJECT_ID    = os.getenv("GOOGLE_PROJECT_ID")
LOCATION      = os.getenv("GOOGLE_LOCATION")
PROCESSOR_ID  = os.getenv("GOOGLE_PROCESSOR_ID")

PAGE_LIMIT = 14  # stay under the 30 page limit


def get_client():
    opts = ClientOptions(api_endpoint=f"{LOCATION}-documentai.googleapis.com")
    return documentai.DocumentProcessorServiceClient(client_options=opts)


def split_pdf(content: bytes, page_limit: int = PAGE_LIMIT) -> list[bytes]:
    """Split PDF bytes into chunks of page_limit pages."""
    import pypdf
    reader = pypdf.PdfReader(io.BytesIO(content))
    total_pages = len(reader.pages)
    chunks = []
    for start in range(0, total_pages, page_limit):
        writer = pypdf.PdfWriter()
        for i in range(start, min(start + page_limit, total_pages)):
            writer.add_page(reader.pages[i])
        buf = io.BytesIO()
        writer.write(buf)
        chunks.append(buf.getvalue())
    return chunks


def ocr_chunk(client, name: str, content: bytes, mime_type: str, page_offset: int) -> list[dict]:
    """OCR a single chunk and return pages with corrected page numbers."""
    raw_doc = documentai.RawDocument(content=content, mime_type=mime_type)
    request = documentai.ProcessRequest(name=name, raw_document=raw_doc)
    result = client.process_document(request=request)
    doc = result.document

    pages = []
    for i, page in enumerate(doc.pages):
        page_text_parts = []
        for block in page.blocks:
            seg = block.layout.text_anchor.text_segments
            for s in seg:
                start = int(s.start_index) if s.start_index else 0
                end   = int(s.end_index)
                page_text_parts.append(doc.text[start:end])
        page_text = "\n".join(page_text_parts).strip()
        if page_text:
            pages.append({"page_num": page_offset + i + 1, "text": page_text})
    return pages


def ocr_file(file_path: str) -> list[dict]:
    client = get_client()
    name = client.processor_path(PROJECT_ID, LOCATION, PROCESSOR_ID)

    ext = os.path.splitext(file_path)[1].lower()
    mime_map = {
        ".pdf":  "application/pdf",
        ".png":  "image/png",
        ".jpg":  "image/jpeg",
        ".jpeg": "image/jpeg",
        ".tiff": "image/tiff",
        ".tif":  "image/tiff",
        ".gif":  "image/gif",
        ".bmp":  "image/bmp",
        ".webp": "image/webp",
    }
    mime_type = mime_map.get(ext, "application/pdf")

    with open(file_path, "rb") as f:
        content = f.read()

    # Only PDFs need splitting
    if mime_type == "application/pdf":
        chunks = split_pdf(content)
        print(f"  Split into {len(chunks)} chunk(s) of up to {PAGE_LIMIT} pages")
    else:
        chunks = [content]

    all_pages = []
    for chunk_idx, chunk in enumerate(chunks):
        page_offset = chunk_idx * PAGE_LIMIT
        print(f"  OCR chunk {chunk_idx + 1}/{len(chunks)}...")
        pages = ocr_chunk(client, name, chunk, mime_type, page_offset)
        all_pages.extend(pages)

    return all_pages