"""Dataset writers for JSONL and Parquet output formats."""

from __future__ import annotations

import json
from collections.abc import Iterable
from pathlib import Path
from typing import Any


def write_jsonl(data: Iterable[dict[str, Any]], path: str | Path) -> int:
    """Write an iterable of dicts to a JSONL file.

    Returns:
        Number of rows written.
    """
    count = 0
    with open(path, "w") as f:
        for row in data:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")
            count += 1
    return count


def write_parquet(data: Iterable[dict[str, Any]], path: str | Path, batch_size: int = 10_000) -> int:
    """Write an iterable of dicts to a Parquet file.

    Returns:
        Number of rows written.
    """
    import pyarrow as pa
    import pyarrow.parquet as pq

    writer = None
    count = 0
    batch: list[dict[str, Any]] = []

    for row in data:
        batch.append(row)
        count += 1

        if len(batch) >= batch_size:
            table = pa.Table.from_pylist(batch)
            if writer is None:
                writer = pq.ParquetWriter(str(path), table.schema)
            writer.write_table(table)
            batch = []

    if batch:
        table = pa.Table.from_pylist(batch)
        if writer is None:
            writer = pq.ParquetWriter(str(path), table.schema)
        writer.write_table(table)

    if writer is not None:
        writer.close()

    return count
