from __future__ import annotations

import csv
import io
import json
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol


class RecordSink(Protocol):
    def on_record(self, *, key: str, record: Any) -> None: ...
    def close(self) -> None: ...


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _normalise_record(record: Any) -> Any:
    """Convert raw record to a JSON-friendly value."""
    if isinstance(record, bytes):
        return record.decode("utf-8", errors="replace")
    v = getattr(record, "list", None)
    if v is not None:
        return v
    if not isinstance(record, (dict, list, str, int, float, bool)) and record is not None:
        return str(record)
    return record


# ---------------------------------------------------------------------------
# Count-only
# ---------------------------------------------------------------------------

@dataclass
class CountSink:
    count: int = 0

    def on_record(self, *, key: str, record: Any) -> None:
        self.count += 1

    def close(self) -> None:
        return


# ---------------------------------------------------------------------------
# In-memory collector
# ---------------------------------------------------------------------------

class CollectorSink:
    """Accumulates records in memory so they can be returned to the caller."""

    def __init__(self) -> None:
        self.count: int = 0
        self.data: dict[str, Any] = {}

    def on_record(self, *, key: str, record: Any) -> None:
        self.data[key] = record
        self.count += 1

    def close(self) -> None:
        return


# ---------------------------------------------------------------------------
# NDJSON / JSONL  (they're the same format)
# ---------------------------------------------------------------------------

class NDJSONSink:
    def __init__(self, path: Path) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        self._fp = path.open("w", encoding="utf-8", buffering=1)  # line-buffered
        self.count = 0

    def on_record(self, *, key: str, record: Any) -> None:
        record_out = _normalise_record(record)
        self._fp.write(json.dumps({"key": key, "record": record_out}) + "\n")
        self.count += 1

    def close(self) -> None:
        self._fp.close()


# ---------------------------------------------------------------------------
# CSV
# ---------------------------------------------------------------------------

class CSVSink:
    """Write records as CSV rows.

    - If the record is a dict, columns are ``key`` + the dict's keys.
    - If the record is a string (dynamic array), columns are ``key, record``.
    """

    def __init__(self, path: Path) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        self._fp = path.open("w", encoding="utf-8", newline="")
        self._writer: csv.DictWriter | csv.writer | None = None
        self._is_dict = False
        self.count = 0

    def on_record(self, *, key: str, record: Any) -> None:
        record_out = _normalise_record(record)

        if self._writer is None:
            if isinstance(record_out, dict):
                self._is_dict = True
                fieldnames = ["key"] + [k for k in record_out if k != "key"]
                self._writer = csv.DictWriter(self._fp, fieldnames=fieldnames)
                self._writer.writeheader()
            else:
                self._is_dict = False
                w = csv.writer(self._fp)
                w.writerow(["key", "record"])
                self._writer = w

        if self._is_dict and isinstance(record_out, dict):
            row = {"key": key, **record_out}
            self._writer.writerow(row)  # type: ignore[union-attr]
        else:
            self._writer.writerow([key, record_out])  # type: ignore[union-attr]
        self.count += 1

    def close(self) -> None:
        self._fp.close()


# ---------------------------------------------------------------------------
# Parquet  (requires pyarrow; graceful error if missing)
# ---------------------------------------------------------------------------

class ParquetSink:
    """Buffer records in memory and flush to a Parquet file on close.

    Requires ``pyarrow``. Install with ``pip install pyarrow``.
    """

    def __init__(self, path: Path) -> None:
        try:
            import pyarrow  # noqa: F401
        except ImportError as exc:
            raise ImportError(
                "ParquetSink requires 'pyarrow'. Install it: pip install pyarrow"
            ) from exc
        path.parent.mkdir(parents=True, exist_ok=True)
        self._path = path
        self._keys: list[str] = []
        self._records: list[Any] = []
        self.count = 0

    def on_record(self, *, key: str, record: Any) -> None:
        record_out = _normalise_record(record)
        self._keys.append(key)
        self._records.append(record_out)
        self.count += 1

    def close(self) -> None:
        import pyarrow as pa
        import pyarrow.parquet as pq

        if not self._keys:
            return

        # If records are dicts, expand into columns. Otherwise, key + record.
        first = self._records[0]
        if isinstance(first, dict):
            columns: dict[str, list[Any]] = {"key": self._keys}
            for field in first:
                columns.setdefault(field, [])
            for rec in self._records:
                if isinstance(rec, dict):
                    for field in columns:
                        if field == "key":
                            continue
                        columns[field].append(rec.get(field))
                else:
                    for field in columns:
                        if field == "key":
                            continue
                        columns[field].append(None)
            table = pa.table(columns)
        else:
            table = pa.table({"key": self._keys, "record": [str(r) for r in self._records]})

        pq.write_table(table, str(self._path))


# ---------------------------------------------------------------------------
# Streaming sinks  (write to a file-like object, typically stdout)
# ---------------------------------------------------------------------------

class StreamJSONSink:
    """Stream NDJSON/JSONL to a file-like object (e.g. stdout)."""

    def __init__(self, fp: Any = None) -> None:
        self._fp = fp or sys.stdout
        self.count = 0

    def on_record(self, *, key: str, record: Any) -> None:
        record_out = _normalise_record(record)
        self._fp.write(json.dumps({"key": key, "record": record_out}) + "\n")
        self._fp.flush()
        self.count += 1

    def close(self) -> None:
        pass  # don't close stdout


class StreamCSVSink:
    """Stream CSV rows to a file-like object (e.g. stdout)."""

    def __init__(self, fp: Any = None) -> None:
        self._fp = fp or sys.stdout
        self._writer: csv.DictWriter | csv.writer | None = None
        self._is_dict = False
        self.count = 0

    def on_record(self, *, key: str, record: Any) -> None:
        record_out = _normalise_record(record)

        if self._writer is None:
            if isinstance(record_out, dict):
                self._is_dict = True
                fieldnames = ["key"] + [k for k in record_out if k != "key"]
                self._writer = csv.DictWriter(self._fp, fieldnames=fieldnames)
                self._writer.writeheader()
            else:
                self._is_dict = False
                w = csv.writer(self._fp)
                w.writerow(["key", "record"])
                self._writer = w

        if self._is_dict and isinstance(record_out, dict):
            row = {"key": key, **record_out}
            self._writer.writerow(row)  # type: ignore[union-attr]
        else:
            self._writer.writerow([key, record_out])  # type: ignore[union-attr]
        self._fp.flush()
        self.count += 1

    def close(self) -> None:
        pass  # don't close stdout
