from typing import List, Any
import os
from ..core.base import Sink

try:
    import pyarrow as pa
    import pyarrow.parquet as pq
except Exception:
    pa = None
    pq = None


class ParquetSink(Sink):
    def __init__(self, path: str) -> None:
        self.path = path
        self._writer = None

    def _ensure_dir(self) -> None:
        d = os.path.dirname(self.path)
        if d and not os.path.exists(d):
            os.makedirs(d, exist_ok=True)

    def write_batch(self, batch: List[Any]) -> None:
        if not batch:
            return
        if pa is None or pq is None:
            raise RuntimeError("pyarrow not available")
        self._ensure_dir()
        table = pa.Table.from_pylist(batch)
        if self._writer is None:
            self._writer = pq.ParquetWriter(self.path, table.schema)
        self._writer.write_table(table)

    def close(self) -> None:
        if self._writer is not None:
            self._writer.close()
            self._writer = None
