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

import logging
import os
import time
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._last_flush: dict[str, datetime] = {}
        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] = []
            self._last_flush[prefix] = datetime.now()

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

        # Extract measurement date/time for partitioning and ID generation
        m_dt = record.get("measurement_datetime")
        m_date_str = str(record.get("measurement_date") or record.get("date") or "")
        m_time_str = str(record.get("measurement_time") or record.get("time") or "")

        # Try to determine partition date
        partition_date = None
        if isinstance(m_dt, datetime):
            partition_date = m_dt.date()
        elif len(m_date_str) == 6:
            try:
                # NMEA date is MMDDYY or YYMMDD
                try:
                    partition_date = datetime.strptime(m_date_str, "%m%d%y").date()
                except ValueError:
                    partition_date = datetime.strptime(m_date_str, "%y%m%d").date()
            except ValueError:
                pass

        if partition_date is None:
            ts = record.get("received_at")
            partition_date = ts.date() if isinstance(ts, datetime) else datetime.now().date()

        record["_partition_date"] = partition_date

        # Add measurement_id for optimized joins
        if "measurement_id" not in record and len(m_date_str) == 6 and len(m_time_str) == 6:
            try:
                # We want YYMMDDHHMMSS for sorting
                if m_date_str[4:6].isdigit():  # Likely MMDDYY
                    yy, mm, dd = m_date_str[4:6], m_date_str[0:2], m_date_str[2:4]
                else:  # Likely YYMMDD
                    yy, mm, dd = m_date_str[0:2], m_date_str[2:4], m_date_str[4:6]
                record["measurement_id"] = int(f"{yy}{mm}{dd}{m_time_str}")
            except (ValueError, TypeError):
                pass

        self._buffers[prefix].append(record)

        if len(self._buffers[prefix]) >= self.buffer_size:
            self.flush(prefix)
        else:
            # Check for stale buffers occasionally even if not full
            # We don't want to check on every single record for high-freq ones,
            # but for low-freq it's fine.
            self.flush_stale(max_age_seconds=300)  # 5 minutes default

    def flush_stale(self, max_age_seconds: int = 300) -> None:
        """Flush buffers that haven't been flushed in a while.

        Args:
            max_age_seconds: Threshold for flushing stale buffers.
        """
        now = datetime.now()
        for prefix in list(self._buffers.keys()):
            if not self._buffers[prefix]:
                continue

            last_flush = self._last_flush.get(prefix)
            if last_flush is None or (now - last_flush).total_seconds() > max_age_seconds:
                logger.info(f"Time-based flush for {prefix} ({len(self._buffers[prefix])} records)")
                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:
                    # Use the pre-calculated partition date
                    date_val = rec.pop("_partition_date", None)
                    if date_val is None:
                        ts = rec.get("received_at")
                        date_val = ts.date() if isinstance(ts, datetime) else datetime.now().date()

                    if date_val not in records_by_date:
                        records_by_date[date_val] = []
                    records_by_date[date_val].append(rec)

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

                self._buffers[p] = []
                self._last_flush[p] = datetime.now()
            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.

        Ensures maximum 1 file per day per prefix by appending to existing file if it exists.
        Also compacts legacy files ({prefix}_*.parquet) into the single daily 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}.parquet (one per daily partition)
        final_filename = f"{prefix}.parquet"
        temp_filename = f"{prefix}.parquet.writing"

        final_path = partition_dir / final_filename
        temp_path = partition_dir / temp_filename

        # Detect legacy files: {prefix}_*.parquet (but not the final one)
        legacy_files = [
            f
            for f in partition_dir.glob(f"{prefix}_*.parquet")
            if f.name != final_filename and not f.name.endswith(".writing")
        ]

        # 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

            # Create dataframe for new records
            new_df = pl.from_dicts(records, infer_schema_length=10000)

            # Collect data from existing and legacy files
            dataframes = []
            if final_path.exists():
                try:
                    dataframes.append(pl.read_parquet(str(final_path)))
                except Exception as e:
                    logger.warning(f"Could not read existing Parquet {final_path}: {e}")

            for legacy_path in legacy_files:
                try:
                    dataframes.append(pl.read_parquet(str(legacy_path)))
                except Exception as e:
                    logger.warning(f"Could not read legacy Parquet {legacy_path}: {e}")

            dataframes.append(new_df)

            # Align types to handle schema conflicts (e.g., Int64 vs Float64)
            if dataframes:
                # Get common columns
                all_cols = set(new_df.columns)
                for df_existing in dataframes[:-1]:  # new_df is at the end
                    all_cols.update(df_existing.columns)

                # For all columns in new_df, align with existing dataframes if types differ
                for col in new_df.columns:
                    for i, df_existing in enumerate(dataframes[:-1]):
                        if col in df_existing.columns:
                            existing_dtype = df_existing[col].dtype
                            new_dtype = new_df[col].dtype
                            if new_dtype != existing_dtype:
                                try:
                                    # Handle Null types on either side:
                                    # Polars cannot cast to/from Null, so rebuild
                                    # the Null-typed column with the real type.
                                    if "Null" in str(new_dtype):
                                        # New data is Null → adopt existing type
                                        new_df = new_df.with_columns(
                                            pl.lit(None, dtype=existing_dtype).alias(col)
                                        )
                                    elif "Null" in str(existing_dtype):
                                        # Existing file is Null → adopt new type
                                        dataframes[i] = df_existing.with_columns(
                                            pl.lit(None, dtype=new_dtype).alias(col)
                                        )
                                    # If target is float-ish, promote to float
                                    elif "Float" in str(existing_dtype):
                                        new_df = new_df.with_columns(
                                            pl.col(col).cast(pl.Float64)
                                        )  # pragma: no cover
                                    else:
                                        new_df = new_df.with_columns(
                                            pl.col(col).cast(existing_dtype)
                                        )
                                except Exception as e:
                                    logger.warning(
                                        f"Failed to align column {col} type from "
                                        f"{new_df[col].dtype} to {existing_dtype}: {e}"
                                    )
                                break
                            break

                # Update the reference in the dataframes list to use the aligned version
                dataframes[-1] = new_df

            # Use diagonal union to handle schema evolutions (new columns)
            df = pl.concat(dataframes, how="diagonal")

            # Backfill measurement_id for legacy records if missing
            if "measurement_id" in df.columns:
                # Fill nulls in measurement_id if date and time are present
                # This is a bit complex in polars expressions, so we do it if possible
                if "measurement_date" in df.columns and "measurement_time" in df.columns:
                    # Only try if we have the necessary strings
                    df = df.with_columns(
                        pl.when(pl.col("measurement_id").is_null())
                        .then(
                            # YYMMDDHHMMSS
                            (
                                pl.col("measurement_date").str.slice(4, 2)
                                + pl.col("measurement_date").str.slice(0, 2)
                                + pl.col("measurement_date").str.slice(2, 2)
                                + pl.col("measurement_time")
                            ).cast(pl.UInt64)
                        )
                        .otherwise(pl.col("measurement_id"))
                        .alias("measurement_id")
                    )

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

            # Atomic replace to final path (atomic on POSIX systems, replaces on Windows if exists)
            # On Windows, this can fail if the file is being read by
            # another process (e.g. DuckDB view)
            max_retries = 5
            for attempt in range(max_retries):
                try:
                    if os.path.exists(final_path):
                        os.replace(temp_path, final_path)
                    else:
                        os.rename(temp_path, final_path)
                    break
                except OSError as e:
                    if attempt == max_retries - 1:
                        logger.error(
                            f"Failed to replace Parquet file "
                            f"{final_path} after {max_retries} "
                            f"attempts: {e}"
                        )
                        raise
                    time.sleep(0.1 * (2**attempt))  # Exponential backoff

            # Clean up legacy files after successful write
            for legacy_path in legacy_files:
                try:
                    legacy_path.unlink()
                except Exception as e:
                    logger.warning(f"Could not delete legacy file {legacy_path}: {e}")

            logger.debug(f"Wrote {len(records)} records (total {len(df)}) 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()
        try:
            self._conn.close()
        except Exception as e:
            logger.error(f"Error closing DuckDB connection in ParquetWriter: {e}")
