"""Parquet writer for efficient storage of structured ADCP records."""

import contextlib
import logging
import os
from datetime import date, datetime
from pathlib import Path
from typing import Any

import duckdb

logger = logging.getLogger(__name__)


class ParquetWriter:
    """Writes structured records to Parquet files with daily partitioning.

    Uses DuckDB as the engine for efficient Parquet generation.
    """

    def __init__(self, base_path: str, buffer_size: int = 100):
        """Initialize Parquet writer.

        Args:
            base_path: Base directory for the "DuckLake" storage
            buffer_size: Number of records to buffer before flushing to disk

        """
        self.base_path = Path(base_path) / "parquet"
        self.buffer_size = buffer_size
        self._buffers: dict[str, list[dict[str, Any]]] = {}
        self._conn = duckdb.connect(database=":memory:")
        self._ensure_base_path()

    def _ensure_base_path(self) -> None:
        """Ensure base directory for Parquet files exists."""
        os.makedirs(self.base_path, exist_ok=True)

    def _get_partition_path(self, prefix: str, record_date: date) -> Path:
        """Get the partitioned directory path for a record type and date."""
        # Partitioning by record type and then by date
        # Format: base/prefix/date=YYYY-MM-DD/
        partition_dir = self.base_path / prefix / f"date={record_date.isoformat()}"
        os.makedirs(partition_dir, exist_ok=True)
        return partition_dir

    def write_record(self, prefix: str, record: dict[str, Any]) -> None:
        """Buffer a record for writing.

        Args:
            prefix: Record type prefix (e.g., 'PNORS', 'PNORC')
            record: Dictionary of data to store

        """
        if prefix not in self._buffers:
            self._buffers[prefix] = []

        # Add timestamp if not present
        if "received_at" not in record:
            record["received_at"] = datetime.now()

        # Add measurement_id for optimized joins if date and time are present
        if "measurement_date" in record and "measurement_time" in record:
            try:
                date_str = str(record["measurement_date"])
                time_str = str(record["measurement_time"])
                if len(date_str) == 6 and len(time_str) == 6:
                    record["measurement_id"] = int(date_str + time_str)
            except (ValueError, TypeError):
                pass

        self._buffers[prefix].append(record)

        if len(self._buffers[prefix]) >= self.buffer_size:
            self.flush(prefix)

    def flush(self, prefix: str | None = None) -> None:
        """Flush buffered records to Parquet files.

        Args:
            prefix: If specified, only flush that prefix. Otherwise flush all.

        """
        prefixes = [prefix] if prefix else list(self._buffers.keys())

        for p in prefixes:
            buffer = self._buffers.get(p)
            if not buffer:
                continue

            try:
                # Group by date for partitioning
                records_by_date: dict[date, list[dict[str, Any]]] = {}
                for rec in buffer:
                    ts = rec.get("received_at")
                    date = ts.date() if isinstance(ts, datetime) else datetime.now().date()
                    if date not in records_by_date:
                        records_by_date[date] = []
                    records_by_date[date].append(rec)

                for date, records in records_by_date.items():
                    self._write_to_parquet(p, date, records)

                self._buffers[p] = []
            except Exception as e:
                logger.error(f"Failed to flush Parquet records for {p}: {e}")

    def _write_to_parquet(
        self, prefix: str, record_date: date, records: list[dict[str, Any]]
    ) -> None:
        """Actually write a batch of records to a Parquet file.

        Uses atomic write signaling: writes to a temporary .writing file first,
        then renames to the final .parquet file. This ensures readers never see
        incomplete files.
        """
        partition_dir = self._get_partition_path(prefix, record_date)

        # Filename: {prefix}_{timestamp}.parquet
        timestamp_str = datetime.now().strftime("%H%M%S_%f")
        final_filename = f"{prefix}_{timestamp_str}.parquet"
        temp_filename = f"{prefix}_{timestamp_str}.parquet.writing"

        final_path = partition_dir / final_filename
        temp_path = partition_dir / temp_filename

        # Ensure all records have record_type for consistency
        for r in records:
            if "record_type" not in r:
                r["record_type"] = prefix

        try:
            # Use polars to write Parquet directly - more efficient and no pandas dependency
            import polars as pl

            # Use from_dicts with large infer_schema_length to avoid schema mismatch issues
            df = pl.from_dicts(records, infer_schema_length=10000)

            # Write to temporary file first
            df.write_parquet(str(temp_path))

            # Atomic rename to final path (atomic on POSIX systems)
            os.rename(temp_path, final_path)

            logger.debug(f"Wrote {len(records)} records to {final_path}")
        except Exception as e:
            logger.error(f"Polars Parquet write error: {prefix}: {e}")
            # Clean up temp file if it exists
            if temp_path.exists():
                try:
                    temp_path.unlink()
                except OSError:
                    pass
            raise

    def close(self) -> None:
        """Flush all buffers and close connections."""
        self.flush()
        with contextlib.suppress(Exception):
            self._conn.close()
