"""Extract text and token count per page from a PDF. Supports PyPDF2 and PyMuPDF."""

import os
from io import BytesIO

import litellm


def get_pdf_page_tokens(pdf_path, model: str = None,
                        pdf_parser: str = "PyPDF2") -> list[tuple[str, int]]:
    """Extract text and token count for each page of a PDF.

    Args:
        pdf_path: Path to PDF file, or BytesIO object.
        model: Model name for token counting (passed to litellm.token_counter).
        pdf_parser: Parser backend — 'PyPDF2' or 'PyMuPDF'.

    Returns:
        list[tuple[str, int]]: List of (page_text, token_count) per page.
    """
    if pdf_parser == "PyPDF2":
        import PyPDF2
        pdf_reader = PyPDF2.PdfReader(pdf_path)
        page_list = []
        for page in pdf_reader.pages:
            page_text = page.extract_text() or ""
            token_length = litellm.token_counter(model=model, text=page_text)
            page_list.append((page_text, token_length))
        return page_list

    elif pdf_parser == "PyMuPDF":
        import pymupdf
        if isinstance(pdf_path, BytesIO):
            doc = pymupdf.open(stream=pdf_path, filetype="pdf")
        elif isinstance(pdf_path, str) and os.path.isfile(pdf_path):
            doc = pymupdf.open(pdf_path)
        else:
            raise ValueError(f"Invalid pdf_path: {pdf_path}")
        page_list = []
        for page in doc:
            page_text = page.get_text()
            token_length = litellm.token_counter(model=model, text=page_text)
            page_list.append((page_text, token_length))
        return page_list

    else:
        raise ValueError(f"Unsupported PDF parser: {pdf_parser}. Use 'PyPDF2' or 'PyMuPDF'.")
