import aiosqlite
import json
import os
from typing import Any, Optional, Dict, AsyncIterator
from contextlib import asynccontextmanager
from pspf.state.store import StateStore
from pspf.utils.logging import get_logger

logger = get_logger("SQLiteStateStore")

try:
    import msgpack
    HAS_MSGPACK = True
except ImportError:
    HAS_MSGPACK = False

def serialize_state(value: Any) -> bytes:
    """Serializes state using msgpack if available, otherwise json."""
    if HAS_MSGPACK:
        try:
            return msgpack.packb(value, use_bin_type=True)
        except Exception as e:
            logger.warning(f"msgpack serialization failed: {e}. Falling back to JSON.")
    
    # Fallback or if msgpack not available
    return json.dumps(value).encode("utf-8")

def deserialize_state(data: bytes, key: str, default: Any) -> Any:
    """Deserializes state, attempting msgpack first, then json."""
    if HAS_MSGPACK:
        try:
            return msgpack.unpackb(data, raw=False)
        except Exception:
            pass # Fallthrough to JSON
            
    try:
        return json.loads(data.decode("utf-8"))
    except Exception as e:
        logger.error(f"CORRUPTION DETECTED: Failed to deserialize value for key '{key}': {e}")
        return default

class SQLiteStateStore(StateStore):
    """
    Persistent state store using SQLite.
    Values are stored via msgpack (preferred) or json.
    """
    def __init__(self, path: str, table_name: str = "kv_store") -> None:
        self.path = path
        self.table_name = table_name
        self._db: Optional[aiosqlite.Connection] = None
        self._in_transaction = False

    async def start(self) -> None:
        # Ensure directory exists
        dirname = os.path.dirname(self.path)
        if dirname and not os.path.exists(dirname):
            os.makedirs(dirname)
            
        self._db = await aiosqlite.connect(self.path)
        # Main KV table
        await self._db.execute(
            f"CREATE TABLE IF NOT EXISTS {self.table_name} (key TEXT PRIMARY KEY, value BLOB)"
        )
        # Offset tracking table
        await self._db.execute(
            "CREATE TABLE IF NOT EXISTS pspf_offsets (stream_id TEXT, group_id TEXT, offset TEXT, PRIMARY KEY (stream_id, group_id))"
        )
        await self._db.commit()
        logger.info(f"Opened SQLite State Store at {self.path}")

    @asynccontextmanager
    async def transaction(self) -> AsyncIterator[None]:
        if not self._db: raise RuntimeError("Store not started")
        if self._in_transaction:
            # Nested transactions not supported yet
            yield
            return

        self._in_transaction = True
        try:
            await self._db.execute("BEGIN TRANSACTION")
            yield
            await self._db.commit()
        except Exception as e:
            await self._db.rollback()
            logger.error(f"Transaction failed, rolled back: {e}")
            raise
        finally:
            self._in_transaction = False

    async def stop(self) -> None:
        if self._db:
            await self._db.close()

    async def get(self, key: str, default: Any = None) -> Any:
        if not self._db: raise RuntimeError("Store not started")
        
        async with self._db.execute(f"SELECT value FROM {self.table_name} WHERE key = ?", (key,)) as cursor:
            row = await cursor.fetchone()
            if row:
                return deserialize_state(row[0], key, default)
            return default

    async def put(self, key: str, value: Any, ttl_seconds: Optional[int] = None) -> None:
        if not self._db: raise RuntimeError("Store not started")
        
        data = serialize_state(value)
        await self._db.execute(
            f"INSERT OR REPLACE INTO {self.table_name} (key, value) VALUES (?, ?)",
            (key, data)
        )
        if not self._in_transaction:
            await self._db.commit()

    async def put_batch(self, entries: Dict[str, Any]) -> None:
        if not self._db: raise RuntimeError("Store not started")
        
        # Prepare data
        rows = []
        for k, v in entries.items():
            rows.append((k, serialize_state(v)))
            
        if not rows: return

        await self._db.executemany(
            f"INSERT OR REPLACE INTO {self.table_name} (key, value) VALUES (?, ?)",
            rows
        )
        if not self._in_transaction:
            await self._db.commit()

    async def delete(self, key: str) -> None:
        if not self._db: raise RuntimeError("Store not started")
        await self._db.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,))
        if not self._in_transaction:
            await self._db.commit()

    async def flush(self) -> None:
        if self._db:
            await self._db.commit()

    async def checkpoint(self, stream_id: str, group_id: str, offset: str) -> None:
        """
        Atomically store the offset. 
        Note: If multiple puts were done without commit, this commit will finalize them too.
        """
        if not self._db: raise RuntimeError("Store not started")
        
        await self._db.execute(
            "INSERT OR REPLACE INTO pspf_offsets (stream_id, group_id, offset) VALUES (?, ?, ?)",
            (stream_id, group_id, offset)
        )
        if not self._in_transaction:
            await self._db.commit()

    async def get_checkpoint(self, stream_id: str, group_id: str) -> Optional[str]:
        if not self._db: raise RuntimeError("Store not started")
        
        async with self._db.execute(
            "SELECT offset FROM pspf_offsets WHERE stream_id = ? AND group_id = ?",
            (stream_id, group_id)
        ) as cursor:
            row = await cursor.fetchone()
            return row[0] if row else None
