from __future__ import annotations

from collections.abc import Iterator
from pathlib import Path

import polars as pl
from polars.io.plugins import register_io_source

from . import _aci_rssa
from ._columns import TRACK_SCHEMA


def scan_tracks(
    path: Path | str, default_batch_size: int = 100_000_000
) -> pl.LazyFrame:
    """Scans an RSSA file to return a Polars LazyFrame instead of eagerly reading all
    tracks into memory. This is the recommended way to work with very large (100+ Gb)
    RSSA files: decoding happens in Rust, batch by batch, driven by Polars' query
    engine.

    Parameters
    ----------
    path: Path | str
        The RSSA file to scan.
    default_batch_size: int
        Number of tracks decoded per batch when Polars doesn't request a specific
        batch size. The right value depends on available memory; a good starting
        point is on the order of 10-100 million tracks (roughly 1-8 Gb per batch).
    """
    path = str(path)

    def source_generator(
        with_columns: list[str] | None,
        predicate: pl.Expr | None,
        n_rows: int | None,
        batch_size: int | None,
    ) -> Iterator[pl.DataFrame]:
        effective_batch_size = batch_size or default_batch_size
        if n_rows is not None:
            # Avoid decoding more than what was asked for when only a handful of
            # rows are needed (e.g. `scan_tracks(...).head(20)`).
            effective_batch_size = min(effective_batch_size, max(n_rows, 1))

        reader = _aci_rssa.RssaBatchReader(path, effective_batch_size, with_columns)
        rows_left = n_rows
        for df in reader:
            if rows_left is not None:
                if df.height > rows_left:
                    df = df.head(rows_left)
                rows_left -= df.height
            if predicate is not None:
                df = df.filter(predicate)
            yield df
            if rows_left is not None and rows_left <= 0:
                break

    return register_io_source(io_source=source_generator, schema=TRACK_SCHEMA)
