# State management for vendor bridge integrations.
# Provides persistent storage for cumulative metrics across restarts.

import logging
import sqlite3
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Dict, Optional, Union

logger = logging.getLogger(__name__)


class StateManager(ABC):
    """Abstract base class for state persistence."""

    @abstractmethod
    def connect(self) -> None:
        """Establish connection to state storage."""
        pass

    @abstractmethod
    def load_state(self) -> Dict[str, int]:
        """Load state from storage."""
        pass

    @abstractmethod
    def save_state(self, state: Dict[str, int]) -> None:
        """Save state to storage."""
        pass

    @abstractmethod
    def close(self) -> None:
        """Close connection to state storage."""
        pass


class SQLiteStateManager(StateManager):
    """Manages state persistence using SQLite database.

    Provides atomic, crash-safe storage for cumulative metrics.
    Uses WAL mode for better concurrency and crash recovery.

    Args:
        db_path: Path to SQLite database file. Use None for in-memory storage.

    Example:
        >>> manager = SQLiteStateManager("/var/lib/mca/state.db")
        >>> manager.connect()
        >>> state = manager.load_state()
        >>> state["predictions"] = state.get("predictions", 0) + 100
        >>> manager.save_state(state)
        >>> manager.close()
    """

    def __init__(self, db_path: Optional[str]):
        """Initialize state manager with database path."""
        if db_path is None:
            self.db_path: Union[str, Path] = ":memory:"
            self._conn: Optional[sqlite3.Connection] = None
            logger.info("Using in-memory state persistence (state will be lost on restart)")
        else:
            self.db_path = Path(db_path)
            self._conn = None
            logger.info(f"Using SQLite state persistence: {self.db_path}")

    def connect(self) -> None:
        """Establish connection to the database and create table if needed."""
        try:
            if self.db_path != ":memory:" and isinstance(self.db_path, Path):
                self.db_path.parent.mkdir(parents=True, exist_ok=True)

            self._conn = sqlite3.connect(str(self.db_path), timeout=10, check_same_thread=False)
            self._conn.execute("PRAGMA journal_mode=WAL;")

            # Create table atomically
            with self._conn:
                self._conn.execute("""
                    CREATE TABLE IF NOT EXISTS state (
                        key TEXT PRIMARY KEY,
                        value INTEGER NOT NULL
                    )
                """)

            logger.info(f"SQLite state database initialized at {self.db_path}")
        except sqlite3.DatabaseError as e:
            # Database file is corrupted, delete and recreate
            logger.warning(
                f"Database file corrupted at {self.db_path}: {e}. Deleting and recreating."
            )
            if self._conn:
                try:
                    self._conn.close()
                except Exception:
                    pass
                self._conn = None

            # Delete corrupted file and retry
            if (
                self.db_path != ":memory:"
                and isinstance(self.db_path, Path)
                and self.db_path.exists()
            ):
                self.db_path.unlink()

            # Retry connection with fresh database
            self._conn = sqlite3.connect(str(self.db_path), timeout=10, check_same_thread=False)
            self._conn.execute("PRAGMA journal_mode=WAL;")
            with self._conn:
                self._conn.execute("""
                    CREATE TABLE IF NOT EXISTS state (
                        key TEXT PRIMARY KEY,
                        value INTEGER NOT NULL
                    )
                """)
            logger.info(f"SQLite state database recreated at {self.db_path}")
        except sqlite3.Error as e:
            logger.error(f"Failed to connect or initialize SQLite database at {self.db_path}: {e}")
            raise

    def load_state(self) -> Dict[str, int]:
        """Load state from the database.

        Returns:
            Dict mapping state keys to integer values.
        """
        if not self._conn:
            self.connect()

        try:
            with self._conn:
                cursor = self._conn.cursor()
                cursor.execute("SELECT key, value FROM state")
                rows = cursor.fetchall()
                state = {key: value for key, value in rows}

                logger.info(f"Loaded state from {self.db_path}: {state}")
                return state
        except sqlite3.Error as e:
            logger.warning(
                f"Failed to load state from {self.db_path}: {e}. Starting with empty state."
            )
            return {}

    def save_state(self, state: Dict[str, int]) -> None:
        """Persist state to the database atomically.

        Args:
            state: Dictionary mapping state keys to integer values.
        """
        if not self._conn:
            logger.error("Cannot save state, no database connection.")
            return

        try:
            with self._conn:
                for key, value in state.items():
                    self._conn.execute(
                        "INSERT OR REPLACE INTO state (key, value) VALUES (?, ?)", (key, value)
                    )
        except sqlite3.Error as e:
            logger.error(f"Failed to save state to {self.db_path}: {e}")

    def close(self) -> None:
        """Close the database connection."""
        if self._conn:
            self._conn.close()
            self._conn = None
            logger.info(f"Closed connection to state database: {self.db_path}")


class InMemoryStateManager(StateManager):
    """In-memory state manager for testing and development.

    WARNING: State is lost on process restart. Not suitable for production.
    """

    def __init__(self) -> None:
        """Initialize in-memory state manager."""
        self._state: Dict[str, int] = {}
        logger.warning("Using in-memory state manager. State will be lost on restart.")

    def connect(self) -> None:
        """No-op for in-memory storage."""
        pass

    def load_state(self) -> Dict[str, int]:
        """Load state from memory."""
        return self._state.copy()

    def save_state(self, state: Dict[str, int]) -> None:
        """Save state to memory."""
        self._state.update(state)

    def close(self) -> None:
        """No-op for in-memory storage."""
        pass
