#!/usr/bin/env python3
"""Persistent state manager backed by SQLite in WAL mode.

A dedicated writer thread coalesces writes with a 500ms flush interval
to reduce SD card wear.  Reads are served directly (WAL allows concurrent
readers).  Module status and generic key-value data are stored in separate
tables.
"""

import os
import time
import json
import sqlite3
import threading
import logging
from typing import Any, Optional
from pathlib import Path

logger = logging.getLogger(__name__)

_DEFAULT_DB = os.path.join(
    os.path.expanduser("~"), ".implant", "state.db"
)


class StateManager:
    """SQLite-backed persistent state with coalesced writes.

    Thread-safe.  The write queue is drained every 500ms by a background
    thread that batches all pending operations into a single transaction.
    """

    FLUSH_INTERVAL = 0.5  # seconds

    def __init__(self, db_path: str = _DEFAULT_DB):
        self._db_path = db_path
        Path(db_path).parent.mkdir(parents=True, exist_ok=True)

        # Writer thread owns this connection
        self._write_conn: Optional[sqlite3.Connection] = None
        # Per-thread read connections via _local
        self._local = threading.local()

        self._write_queue: list[tuple] = []  # (sql, params) tuples
        self._queue_lock = threading.Lock()
        self._writer_thread: Optional[threading.Thread] = None
        self._running = False

        self._init_db()

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

    def start(self) -> None:
        """Start the background writer thread."""
        if self._running:
            return
        self._running = True
        self._writer_thread = threading.Thread(
            target=self._writer_loop, daemon=True, name="sensor-state-writer"
        )
        self._writer_thread.start()
        logger.info("StateManager writer started (db=%s)", self._db_path)

    def stop(self, secure_wipe: bool = False) -> None:
        """Flush remaining writes and stop the writer thread.
        
        Args:
            secure_wipe: If True, securely overwrite and delete the state database.
        """
        if not self._running:
            return
        self._running = False
        if self._writer_thread and self._writer_thread.is_alive():
            self._writer_thread.join(timeout=5.0)
        # Final flush
        self._flush()
        if self._write_conn:
            self._write_conn.close()
            self._write_conn = None
        
        # Secure wipe if requested
        if secure_wipe:
            from utils.crypto import secure_wipe_file
            # SQLite creates multiple files in WAL mode: db, db-wal, db-shm
            for suffix in ["", "-wal", "-shm"]:
                filepath = self._db_path + suffix
                if os.path.isfile(filepath):
                    try:
                        if secure_wipe_file(filepath, passes=3):
                            logger.info("Securely wiped %s", filepath)
                        else:
                            logger.warning("Secure wipe failed for %s (used regular delete)", filepath)
                    except Exception as e:
                        logger.error("Failed to wipe %s: %s", filepath, e)
        
        logger.info("StateManager stopped")

    # ------------------------------------------------------------------
    # Schema init
    # ------------------------------------------------------------------

    def _init_db(self) -> None:
        conn = sqlite3.connect(self._db_path)
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA synchronous=NORMAL")
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS kv_store (
                module  TEXT NOT NULL,
                key     TEXT NOT NULL,
                value   TEXT,
                updated REAL NOT NULL,
                PRIMARY KEY (module, key)
            );

            CREATE TABLE IF NOT EXISTS module_status (
                module  TEXT PRIMARY KEY,
                status  TEXT NOT NULL DEFAULT 'stopped',
                pid     INTEGER,
                started REAL,
                extra   TEXT,
                updated REAL NOT NULL
            );
        """)
        conn.commit()
        conn.close()

    def _get_read_conn(self) -> sqlite3.Connection:
        """Return a per-thread read connection."""
        conn = getattr(self._local, "conn", None)
        if conn is None:
            conn = sqlite3.connect(self._db_path, check_same_thread=False)
            conn.execute("PRAGMA journal_mode=WAL")
            conn.execute("PRAGMA synchronous=NORMAL")
            conn.row_factory = sqlite3.Row
            self._local.conn = conn
        return conn

    def _get_write_conn(self) -> sqlite3.Connection:
        if self._write_conn is None:
            self._write_conn = sqlite3.connect(self._db_path, check_same_thread=False, timeout=30)
            self._write_conn.execute("PRAGMA journal_mode=WAL")
            self._write_conn.execute("PRAGMA synchronous=NORMAL")
            self._write_conn.execute("PRAGMA busy_timeout=30000")
        return self._write_conn

    # ------------------------------------------------------------------
    # Key-value store
    # ------------------------------------------------------------------

    def set(self, module: str, key: str, value: Any) -> None:
        """Set a key-value pair (coalesced write)."""
        val_str = json.dumps(value) if not isinstance(value, str) else value
        sql = """INSERT INTO kv_store (module, key, value, updated)
                 VALUES (?, ?, ?, ?)
                 ON CONFLICT(module, key)
                 DO UPDATE SET value=excluded.value, updated=excluded.updated"""
        self._enqueue(sql, (module, key, val_str, time.time()))

    def get(self, module: str, key: str, default: Any = None) -> Any:
        """Get a value by module + key (direct read, no queue delay)."""
        conn = self._get_read_conn()
        row = conn.execute(
            "SELECT value FROM kv_store WHERE module=? AND key=?",
            (module, key),
        ).fetchone()
        if row is None:
            return default
        return row[0]

    def get_all(self, module: str) -> dict:
        """Get all key-value pairs for a module."""
        conn = self._get_read_conn()
        rows = conn.execute(
            "SELECT key, value FROM kv_store WHERE module=?", (module,)
        ).fetchall()
        return {r[0]: r[1] for r in rows}

    def delete(self, module: str, key: str) -> None:
        """Delete a key-value pair."""
        self._enqueue("DELETE FROM kv_store WHERE module=? AND key=?", (module, key))

    # ------------------------------------------------------------------
    # Module status
    # ------------------------------------------------------------------

    def set_module_status(self, module: str, status: str, pid: int = None,
                          extra: dict = None) -> None:
        """Update module runtime status."""
        extra_str = json.dumps(extra) if extra else None
        sql = """INSERT INTO module_status (module, status, pid, started, extra, updated)
                 VALUES (?, ?, ?, ?, ?, ?)
                 ON CONFLICT(module)
                 DO UPDATE SET status=excluded.status, pid=excluded.pid,
                               started=CASE WHEN excluded.status='running'
                                            THEN excluded.started
                                            ELSE module_status.started END,
                               extra=excluded.extra, updated=excluded.updated"""
        started = time.time() if status == "running" else None
        self._enqueue(sql, (module, status, pid, started, extra_str, time.time()))

    def get_module_status(self, module: str) -> dict:
        """Get module runtime status."""
        conn = self._get_read_conn()
        row = conn.execute(
            "SELECT status, pid, started, extra, updated FROM module_status WHERE module=?",
            (module,),
        ).fetchone()
        if row is None:
            return {"status": "unknown", "pid": None, "started": None, "extra": None}
        return {
            "status": row[0],
            "pid": row[1],
            "started": row[2],
            "extra": json.loads(row[3]) if row[3] else None,
            "updated": row[4],
        }

    def get_all_module_status(self) -> dict:
        """Get status for all modules."""
        conn = self._get_read_conn()
        rows = conn.execute(
            "SELECT module, status, pid, started, extra, updated FROM module_status"
        ).fetchall()
        return {
            r[0]: {
                "status": r[1], "pid": r[2], "started": r[3],
                "extra": json.loads(r[4]) if r[4] else None,
                "updated": r[5],
            }
            for r in rows
        }

    def reset_module_status(self) -> None:
        """Mark all running modules as stopped.

        Called at engine startup to clear stale PIDs from a previous run.
        Without this, the watchdog picks up old PIDs and reports dead modules
        that were never started in the current session.
        """
        conn = self._get_write_conn()
        try:
            conn.execute("BEGIN IMMEDIATE")
            conn.execute(
                "UPDATE module_status SET status='stopped', pid=NULL, updated=? "
                "WHERE status='running'",
                (time.time(),),
            )
            conn.commit()
        except Exception:
            try:
                conn.rollback()
            except Exception:
                pass

    # ------------------------------------------------------------------
    # Write queue
    # ------------------------------------------------------------------

    def _enqueue(self, sql: str, params: tuple) -> None:
        with self._queue_lock:
            self._write_queue.append((sql, params))

    def _flush(self) -> None:
        """Execute all queued writes in a single transaction."""
        with self._queue_lock:
            batch = list(self._write_queue)
            self._write_queue.clear()

        if not batch:
            return

        conn = self._get_write_conn()
        for attempt in range(6):
            try:
                conn.execute("BEGIN IMMEDIATE")
                for sql, params in batch:
                    conn.execute(sql, params)
                conn.execute("COMMIT")
                return
            except sqlite3.OperationalError as e:
                try:
                    conn.execute("ROLLBACK")
                except Exception:
                    pass
                if attempt < 5:
                    time.sleep(0.5 * (attempt + 1))
                else:
                    logger.error("StateManager flush failed after 6 attempts (%d ops): %s", len(batch), e)
            except Exception:
                try:
                    conn.execute("ROLLBACK")
                except Exception:
                    pass
                logger.exception("StateManager flush failed (%d ops)", len(batch))
                return

    def _writer_loop(self) -> None:
        """Background loop: flush write queue every FLUSH_INTERVAL seconds."""
        while self._running:
            time.sleep(self.FLUSH_INTERVAL)
            try:
                self._flush()
            except Exception:
                logger.exception("Writer loop error")

    # ------------------------------------------------------------------
    # Maintenance
    # ------------------------------------------------------------------

    def checkpoint(self) -> None:
        """Force a WAL checkpoint (periodic call to limit WAL size)."""
        conn = self._get_write_conn()
        conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")

    def flush_sync(self) -> None:
        """Synchronously flush all pending writes (for shutdown paths)."""
        self._flush()
