from __future__ import annotations
import fitz  # pymupdf
from ..models import DocumentContent, TextSegment


def parse_pdf(
    buffer: bytes,
    max_pages: int | None = None,
    pages: tuple[int, int] | None = None,
) -> DocumentContent:
    """Parse a PDF and return its text content as a ``DocumentContent``.

    Parameters
    ----------
    buffer:
        Raw PDF bytes.
    max_pages:
        If set, parse only the first *N* pages.  Ignored when ``pages`` is given.
    pages:
        1-indexed ``(start, end)`` page range (both inclusive).  E.g.
        ``pages=(50, 100)`` extracts pages 50–100.  Useful for jumping directly
        to a known section of a large tariff PDF.
    """
    doc = fitz.open(stream=buffer, filetype="pdf")
    total_page_count = doc.page_count

    # Resolve page range (fitz uses 0-indexed)
    if pages is not None:
        start_idx = max(0, pages[0] - 1)
        end_idx = min(total_page_count, pages[1])  # exclusive upper bound
    else:
        start_idx = 0
        end_idx = total_page_count if max_pages is None else min(max_pages, total_page_count)

    segments: list[TextSegment] = []
    page_dimensions: list[dict] = []
    raw_text = ""

    for page_idx in range(start_idx, end_idx):
        page = doc[page_idx]
        page_number = page_idx + 1

        rect = page.rect
        page_dimensions.append({"width": rect.width, "height": rect.height})

        for block in page.get_text("dict")["blocks"]:
            if block.get("type") != 0:
                continue
            for line in block.get("lines", []):
                line_text = "".join(span.get("text", "") for span in line.get("spans", []))
                trimmed = line_text.strip()
                if not trimmed:
                    continue

                bbox = line.get("bbox") or block.get("bbox")
                segments.append(TextSegment(
                    text=trimmed,
                    charOffset=len(raw_text),
                    pageNumber=page_number,
                    x=bbox[0] if bbox else None,
                    y=bbox[1] if bbox else None,
                    width=(bbox[2] - bbox[0]) if bbox else None,
                    height=(bbox[3] - bbox[1]) if bbox else None,
                ))
                raw_text += trimmed + "\n"

        if page_idx < end_idx - 1:
            raw_text += "\n"

    doc.close()
    return DocumentContent(
        rawText=raw_text.rstrip(),
        segments=segments,
        pageCount=total_page_count,
        pageDimensions=page_dimensions,
    )
