"""Database writer for persisting MQTT telemetry into TimescaleDB.

Messages are enqueued by the MQTT callback thread and flushed to the
database in batches by a dedicated worker thread. This keeps blocking I/O
off the MQTT network loop and amortizes commit overhead over many rows.
"""

import logging
import os
import queue
import threading
import time
from datetime import datetime, timezone
from typing import Any, Optional

import psycopg
from psycopg.types.json import Jsonb

logger = logging.getLogger("mqtt-client")

_INSERT_SQL = """
    INSERT INTO telemetry (time, topic, device_id, payload)
    VALUES (%s, %s, %s, %s)
"""


class DbWriter:
    """Queue-backed batch writer for telemetry rows.

    ``write()`` only enqueues, so it is safe to call from the MQTT network
    loop. A worker thread groups rows into batches and inserts each batch
    in a single transaction; failed batches are retried and re-enqueued
    rather than dropped.
    """

    def __init__(
        self,
        batch_size: Optional[int] = None,
        flush_interval: Optional[float] = None,
        queue_limit: int = 10_000,
    ) -> None:
        if batch_size is None:
            batch_size = int(os.getenv("DB_BATCH_SIZE", "500"))
        if flush_interval is None:
            flush_interval = float(os.getenv("DB_FLUSH_SECONDS", "1.0"))
        self._batch_size = max(1, batch_size)
        self._flush_interval = max(0.05, flush_interval)
        self._queue: queue.Queue[tuple[datetime, str, Any]] = queue.Queue(
            maxsize=queue_limit
        )
        self._stop_event = threading.Event()
        self._drained = threading.Event()
        self.conn: Optional[psycopg.Connection] = None

        self._worker_thread = threading.Thread(
            target=self._worker, name="db-writer", daemon=True
        )
        self._worker_thread.start()

    # -- Producer side (called from the MQTT callback thread) -----------

    def write(self, topic: str, payload: Any) -> None:
        """Enqueue a message for persistence; never blocks the caller."""
        if self._stop_event.is_set():
            logger.error("Writer stopped; dropping message on %s", topic)
            return
        try:
            self._queue.put_nowait((datetime.now(timezone.utc), topic, payload))
        except queue.Full:
            # Make backpressure visible instead of stalling the MQTT loop.
            logger.error(
                "Write queue full (%d pending); dropping message on %s",
                self._queue.maxsize,
                topic,
            )

    # -- Consumer side (worker thread) ------------------------------------

    def _worker(self) -> None:
        while not self._stop_event.is_set() or not self._queue.empty():
            batch = self._collect_batch()
            if batch:
                self._flush_with_retry(batch)
        self._drained.set()
        self._close_connection()

    def _collect_batch(self) -> list[tuple[datetime, str, Any]]:
        """Block for up to one flush interval, then drain up to batch size."""
        batch: list[tuple[datetime, str, Any]] = []
        try:
            batch.append(self._queue.get(timeout=self._flush_interval))
        except queue.Empty:
            return batch
        while len(batch) < self._batch_size:
            try:
                batch.append(self._queue.get_nowait())
            except queue.Empty:
                break
        return batch

    def _flush_with_retry(self, batch: list[tuple[datetime, str, Any]]) -> None:
        for attempt in (1, 2):
            try:
                self._flush(batch)
                return
            except Exception:  # noqa: BLE001 - any DB error is retryable
                logger.exception(
                    "Failed to flush %d rows (attempt %d/2)", len(batch), attempt
                )
                self._close_connection()
                if attempt == 1:
                    time.sleep(2.0)

        if self._stop_event.is_set():
            logger.error("Shutting down; dropping %d unwritten rows", len(batch))
            return
        # Keep the data; it is retried on the next loop iteration.
        logger.error("Re-enqueueing %d rows after failed flush", len(batch))
        for item in batch:
            try:
                self._queue.put_nowait(item)
            except queue.Full:
                logger.error("Write queue full while re-enqueueing; dropping row")

    def _flush(self, batch: list[tuple[datetime, str, Any]]) -> None:
        self.ensure_connected()
        assert self.conn is not None
        rows = [
            (ts, topic, self._device_id(payload), Jsonb(payload))
            for ts, topic, payload in batch
        ]
        with self.conn.cursor() as cur, self.conn.transaction():
            cur.executemany(_INSERT_SQL, rows)
        logger.debug("Flushed %d rows to TimescaleDB", len(rows))

    # -- Connection management ---------------------------------------------

    def connect(self) -> None:
        self.conn = psycopg.connect(
            host=os.getenv("DB_HOST", "timescaledb"),
            port=int(os.getenv("DB_PORT", "5432")),
            dbname=os.getenv("DB_NAME", "telemetry"),
            user=os.getenv("DB_USER", "ingester"),
            password=os.getenv("DB_PASSWORD", "changeme"),
            autocommit=True,
        )
        logger.info("Connected to TimescaleDB")

    def ensure_connected(self) -> None:
        if self.conn is None or self.conn.closed:
            self.connect()
            return
        try:
            # Cheap liveness check: a silently-dead socket (DB restart,
            # network blip) would otherwise only surface on the next write.
            self.conn.execute("SELECT 1").fetchone()
        except psycopg.Error:
            logger.warning("Database connection stale; reconnecting")
            self._close_connection()
            self.connect()

    def _close_connection(self) -> None:
        if self.conn is not None:
            try:
                self.conn.close()
            except Exception:  # noqa: BLE001 - best-effort cleanup
                pass
            self.conn = None

    # -- Lifecycle -----------------------------------------------------------

    def close(self, timeout: float = 15.0) -> None:
        """Stop the worker thread, flushing any queued rows first."""
        if self._stop_event.is_set():
            return
        self._stop_event.set()
        if not self._drained.wait(timeout):
            logger.warning("Timed out flushing write queue; some rows may be lost")
        self._close_connection()

    @staticmethod
    def _device_id(payload: Any) -> Optional[str]:
        if isinstance(payload, dict):
            return payload.get("device_id") or payload.get("deviceId")
        return None