import asyncio
import threading
from typing import Any, Self, Type, AsyncContextManager

import aiosqlite
from pydantic import BaseModel

from .interfaces import (
    IBeaverBlob,
    IBeaverChannel,
    IBeaverDict,
    IBeaverDocuments,
    IBeaverEvents,
    IBeaverGraph,
    IBeaverList,
    IBeaverLock,
    IBeaverLog,
    IBeaverQueue,
    IBeaverSketch,
    IBeaverVectors,
)


from .blobs import AsyncBeaverBlob
from .bridge import BeaverBridge
from .cache import DummyCache
from .channels import AsyncBeaverChannel, PubSubEngine
from .dicts import AsyncBeaverDict
from .docs import AsyncBeaverDocuments
from .events import AsyncBeaverEvents
from .graphs import AsyncBeaverGraph
from .lists import AsyncBeaverList
from .locks import AsyncBeaverLock
from .logs import AsyncBeaverLog
from .manager import AsyncBeaverBase
from .queues import AsyncBeaverQueue
from .sketches import AsyncBeaverSketch
from .vectors import AsyncBeaverVectors


class Transaction:
    """
    A Reentrant Async Context Manager for database transactions.

    If a task already holds the transaction lock (nested @atomic calls),
    this acts as a pass-through. The actual BEGIN/COMMIT only happens
    at the outermost level.
    """

    def __init__(self, db: "AsyncBeaverDB"):
        self.db = db
        self._is_root = False

    async def __aenter__(self):
        current_task = asyncio.current_task()

        # 1. Check Reentrancy: Do we already own the transaction?
        if self.db._tx_owner_task == current_task:
            self._is_root = False
            return self

        # 2. Acquire Lock (Wait for other tasks)
        await self.db._tx_lock.acquire()

        # 3. Mark ownership and start DB transaction
        self.db._tx_owner_task = current_task
        self._is_root = True

        # Using isolation_level=None globally means we MUST start manually
        await self.db.connection.execute("BEGIN IMMEDIATE")
        return self

    async def __aexit__(self, exc_type, exc, tb):
        # Only the root transaction performs the Commit/Rollback and Release
        if not self._is_root:
            return

        try:
            if exc_type:
                await self.db.connection.rollback()
            else:
                await self.db.connection.commit()
        finally:
            # Always clear ownership and release lock
            self.db._tx_owner_task = None
            self.db._tx_lock.release()


class AsyncBeaverDB:
    """
    The Async-First Core Engine of BeaverDB.

    This class manages the single aiosqlite connection and strictly runs
    within an asyncio event loop. It is NOT thread-safe; it is designed
    to be owned by a single thread (the Reactor Thread).
    """

    def __init__(
        self,
        db_path: str,
        /,
        *,
        connection_timeout: float = 30.0,
        cache_timeout: float = 0.0,
        pragma_wal: bool = True,
        pragma_synchronous: bool = False,
        pragma_temp_memory: bool = True,
        pragma_mmap_size: int = 256 * 1024 * 1024,
    ):
        self._db_path = db_path
        self._timeout = connection_timeout
        self._cache_timeout = cache_timeout

        # The Single Source of Truth Connection
        self._connection: aiosqlite.Connection | None = None

        # Locking Primitives
        self._tx_lock = asyncio.Lock()
        self._tx_owner_task: asyncio.Task | None = None  # Track owner for reentrancy

        # Manager Singleton Cache
        self._manager_cache: dict[tuple[type, str], Any] = {}

        # Store pragma settings
        self._pragma_wal = pragma_wal
        self._pragma_synchronous = pragma_synchronous
        self._pragma_temp_memory = pragma_temp_memory
        self._pragma_mmap_size = pragma_mmap_size

        # Pub/Sub Engine
        self._pubsub_engine: PubSubEngine | None = None

    async def pubsub_engine(self) -> PubSubEngine:
        if self._pubsub_engine is None:
            self._pubsub_engine = PubSubEngine(self)
            await self._pubsub_engine.start()

        return self._pubsub_engine

    async def connect(self) -> Self:
        """
        Initializes the async database connection and creates tables.
        Must be awaited before using the DB.
        """
        if self._connection is not None:
            return self

        self._connection = await aiosqlite.connect(
            self._db_path,
            timeout=self._timeout,
            # We will manage transactions manually via .transaction()
            isolation_level=None,
        )
        self._connection.row_factory = aiosqlite.Row

        # Apply Pragmas
        if self._pragma_wal:
            await self._connection.execute("PRAGMA journal_mode = WAL;")

        if self._pragma_synchronous:
            await self._connection.execute("PRAGMA synchronous = FULL;")
        else:
            await self._connection.execute("PRAGMA synchronous = NORMAL;")

        if self._pragma_temp_memory:
            await self._connection.execute("PRAGMA temp_store = MEMORY;")

        if self._pragma_mmap_size > 0:
            await self._connection.execute(
                f"PRAGMA mmap_size = {self._pragma_mmap_size};"
            )

        await self._create_all_tables()
        # await self._check_version()

        return self

    async def close(self):
        """Closes the database connection."""
        if self._connection:
            await self._connection.close()
            self._connection = None

        # Clear cache to allow GC
        self._manager_cache.clear()

    async def __aenter__(self):
        return await self.connect()

    async def __aexit__(self, *args, **kwargs):
        await self.close()

    @property
    def connection(self) -> aiosqlite.Connection:
        """
        Returns the raw aiosqlite connection.
        Raises an error if not connected.
        """
        if self._connection is None:
            raise ConnectionError(
                "AsyncBeaverDB is not connected. Await .connect() first."
            )

        return self._connection

    def transaction(self) -> AsyncContextManager:
        """
        Returns an async context manager for an atomic transaction.
        Use: async with db.transaction(): ...
        """
        return Transaction(self)

    async def _create_all_tables(self):
        """Initializes all required tables with the new __beaver__ naming convention."""
        # Note: We use execute() directly here as these are DDL statements
        # and don't strictly require the transaction lock (sqlite handles DDL locking).

        c = self.connection

        # Blobs
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_blobs__ (
                store_name TEXT NOT NULL,
                key TEXT NOT NULL,
                data BLOB NOT NULL,
                metadata TEXT,
                PRIMARY KEY (store_name, key)
            )
        """
        )

        # Cache Versioning
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_manager_versions__ (
                namespace TEXT PRIMARY KEY,
                version INTEGER NOT NULL DEFAULT 0
            )
        """
        )

        # Dicts
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_dicts__ (
                dict_name TEXT NOT NULL,
                key TEXT NOT NULL,
                value TEXT NOT NULL,
                expires_at REAL,
                PRIMARY KEY (dict_name, key)
            )
        """
        )

        # Lists
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_lists__ (
                list_name TEXT NOT NULL,
                item_order REAL NOT NULL,
                item_value TEXT NOT NULL,
                PRIMARY KEY (list_name, item_order)
            )
        """
        )

        # Locks
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_lock_waiters__ (
                lock_name TEXT NOT NULL,
                waiter_id TEXT NOT NULL,
                requested_at REAL NOT NULL,
                expires_at REAL NOT NULL,
                PRIMARY KEY (lock_name, requested_at)
            )
        """
        )
        await c.execute(
            "CREATE INDEX IF NOT EXISTS idx_lock_expires ON __beaver_lock_waiters__ (lock_name, expires_at)"
        )
        await c.execute(
            "CREATE INDEX IF NOT EXISTS idx_lock_waiter_id ON __beaver_lock_waiters__ (lock_name, waiter_id)"
        )

        # Logs
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_logs__ (
                log_name TEXT NOT NULL,
                timestamp REAL NOT NULL,
                data TEXT NOT NULL,
                PRIMARY KEY (log_name, timestamp)
            )
        """
        )
        await c.execute(
            "CREATE INDEX IF NOT EXISTS idx_logs_timestamp ON __beaver_logs__ (log_name, timestamp)"
        )

        # Priority Queues
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_priority_queues__ (
                queue_name TEXT NOT NULL,
                priority REAL NOT NULL,
                timestamp REAL NOT NULL,
                data TEXT NOT NULL
            )
        """
        )
        await c.execute(
            "CREATE INDEX IF NOT EXISTS idx_priority_queue_order ON __beaver_priority_queues__ (queue_name, priority ASC, timestamp ASC)"
        )

        # PubSub
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_pubsub_log__ (
                timestamp REAL PRIMARY KEY,
                channel_name TEXT NOT NULL,
                message_payload TEXT NOT NULL
            )
        """
        )
        await c.execute(
            "CREATE INDEX IF NOT EXISTS idx_pubsub_channel_timestamp ON __beaver_pubsub_log__ (channel_name, timestamp)"
        )

        # Sketches
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_sketches__ (
                name TEXT PRIMARY KEY,
                type TEXT NOT NULL,
                capacity INTEGER NOT NULL,
                error_rate REAL NOT NULL,
                data BLOB NOT NULL
            )
        """
        )

        # Main Storage
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_documents__ (
                collection TEXT NOT NULL,
                item_id TEXT NOT NULL,
                data TEXT NOT NULL,
                PRIMARY KEY (collection, item_id)
            )
        """
        )

        # FTS Index
        await c.execute(
            """
            CREATE VIRTUAL TABLE IF NOT EXISTS __beaver_fts_index__ USING fts5(
                collection,
                item_id,
                field_path,
                field_content,
                tokenize = 'porter'
            )
        """
        )

        # Fuzzy Index (Trigrams)
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_trigrams__ (
                collection TEXT NOT NULL,
                item_id TEXT NOT NULL,
                trigram TEXT NOT NULL,
                PRIMARY KEY (collection, item_id, trigram)
            )
        """
        )
        await c.execute(
            "CREATE INDEX IF NOT EXISTS idx_trigram_lookup ON __beaver_trigrams__ (collection, trigram)"
        )

        # Edges
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_edges__ (
                collection TEXT NOT NULL,
                source_item_id TEXT NOT NULL,
                target_item_id TEXT NOT NULL,
                label TEXT NOT NULL,
                metadata TEXT,
                PRIMARY KEY (collection, source_item_id, target_item_id, label)
            )
        """
        )

        # Index for reverse lookups (parents)
        await c.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_edges_target
            ON __beaver_edges__ (collection, target_item_id)
        """
        )

        # Vectors (Main Store)
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_vectors__ (
                collection TEXT NOT NULL,
                item_id TEXT NOT NULL,
                vector BLOB NOT NULL,
                metadata TEXT,
                PRIMARY KEY (collection, item_id)
            )
        """
        )

        # --- LSH Strategy Tables ---

        # LSH Config (Hyperplanes per collection)
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_lsh_config__ (
                collection TEXT PRIMARY KEY,
                hyperplanes BLOB NOT NULL
            )
            """
        )

        # LSH Index (Maps buckets to item_ids)
        await c.execute(
            """
            CREATE TABLE IF NOT EXISTS __beaver_lsh_index__ (
                collection TEXT NOT NULL,
                bucket_id INTEGER NOT NULL,
                item_id TEXT NOT NULL,
                PRIMARY KEY (collection, bucket_id, item_id)
            )
            """
        )
        # Optimized index for retrieving items in a bucket
        await c.execute(
            "CREATE INDEX IF NOT EXISTS idx_lsh_lookup ON __beaver_lsh_index__ (collection, bucket_id)"
        )

        await self.connection.commit()

    def singleton[T: AsyncBeaverBase](self, cls: type[T], name, **kwargs) -> T:
        """
        Factory method to get a singleton manager.
        Since this runs on the event loop, no locks are needed.
        """
        cache_key = (cls, name)

        if cache_key not in self._manager_cache:
            # We pass 'self' (AsyncBeaverDB) as the db interface
            instance = cls(name=name, db=self, **kwargs)
            self._manager_cache[cache_key] = instance

        return self._manager_cache[cache_key]

    # --- Factory Methods (Internal) ---
    # These return the raw Async Managers.
    # Note: These manager classes will be refactored in Phase 3.

    def dict[T: BaseModel](
        self, name: str, model: type[T] | None = None, secret: str | None = None
    ) -> AsyncBeaverDict[T]:
        return self.singleton(AsyncBeaverDict, name, model=model, secret=secret)

    def list[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> AsyncBeaverList[T]:
        return self.singleton(AsyncBeaverList, name, model=model)

    def queue[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> AsyncBeaverQueue[T]:
        return self.singleton(AsyncBeaverQueue, name, model=model)

    def docs[T: BaseModel](
        self, name: str, model: Type[T] | None = None
    ) -> AsyncBeaverDocuments[T]:
        return self.singleton(AsyncBeaverDocuments, name, model=model)

    def graphs[T: BaseModel](
        self, name: str, model: Type[T] | None = None
    ) -> AsyncBeaverGraph[T]:
        return self.singleton(AsyncBeaverGraph, name, model=model)

    def channel[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> AsyncBeaverChannel[T]:
        return self.singleton(AsyncBeaverChannel, name, model=model)

    def blob[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> AsyncBeaverBlob[T]:
        return self.singleton(AsyncBeaverBlob, name, model=model)

    def log[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> AsyncBeaverLog[T]:
        return self.singleton(AsyncBeaverLog, name, model=model)

    def lock(
        self, name: str, timeout=None, lock_ttl=60.0, poll_interval=0.1
    ) -> AsyncBeaverLock:
        return AsyncBeaverLock(self, name, timeout, lock_ttl, poll_interval)

    def sketch[T: BaseModel](
        self,
        name: str,
        capacity=1_000_000,
        error_rate=0.01,
        model: type[T] | None = None,
    ) -> AsyncBeaverSketch[T]:
        return self.singleton(
            AsyncBeaverSketch,
            name,
            capacity=capacity,
            error_rate=error_rate,
            model=model,
        )

    def graph[T: BaseModel](
        self, name, model: type[T] | None = None
    ) -> AsyncBeaverGraph[T]:
        return self.singleton(
            AsyncBeaverGraph,
            name,
            model=model,
        )

    def events[T: BaseModel](
        self, name, model: type[T] | None = None
    ) -> AsyncBeaverEvents:
        return self.singleton(
            AsyncBeaverEvents,
            name,
            model=model,
        )

    def vectors[T: BaseModel](
        self, name, model: type[T] | None = None
    ) -> AsyncBeaverVectors:
        return self.singleton(AsyncBeaverVectors, name, model=model)

    def cache(self, key: str = "global"):
        # Temporary stub: Caching will be revisited
        return DummyCache.singleton()


class BeaverDB:
    """
    The Synchronous Facade (Portal).

    This class starts a background thread with an asyncio loop and
    proxies all requests to the AsyncBeaverDB engine via BeaverBridge.
    """

    def __init__(self, db_path: str, /, **kwargs):
        # 1. Start the Reactor Thread
        self._loop = asyncio.new_event_loop()
        self._thread = threading.Thread(
            target=self._loop.run_forever, daemon=True, name="BeaverDB-Reactor"
        )
        self._thread.start()

        # 2. Initialize the Engine on the Reactor Thread
        async def init_engine():
            db = AsyncBeaverDB(db_path, **kwargs)
            await db.connect()
            return db

        future = asyncio.run_coroutine_threadsafe(init_engine(), self._loop)
        self._async_db = future.result()
        self._closed = False

    def close(self):
        """Shuts down the reactor thread and closes the DB."""
        if self._closed:
            return

        async def shutdown():
            await self._async_db.close()

        future = asyncio.run_coroutine_threadsafe(shutdown(), self._loop)
        future.result()

        self._loop.call_soon_threadsafe(self._loop.stop)
        self._thread.join(timeout=1.0)
        self._closed = True

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def _get_manager(self, method_name: str, *args, **kwargs) -> Any:
        """
        Helper to invoke a factory method on the Async Engine and wrap the result.
        Executing on the loop ensures the singleton cache is accessed safely.
        """

        async def factory_call():
            method = getattr(self._async_db, method_name)
            return method(*args, **kwargs)

        future = asyncio.run_coroutine_threadsafe(factory_call(), self._loop)
        async_manager = future.result()

        # Wrap the Async Manager in the Bridge
        return BeaverBridge(async_manager, self._loop)

    # --- Public API (Proxies) ---

    def dict[T: BaseModel](
        self, name: str, model: type[T] | None = None, secret: str | None = None
    ) -> IBeaverDict[T]:
        return self._get_manager("dict", name, model, secret)

    def list[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverList[T]:
        return self._get_manager("list", name, model)

    def queue[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverQueue[T]:
        return self._get_manager("queue", name, model)

    def docs[T: BaseModel](
        self, name: str, model: Type[T] | None = None
    ) -> IBeaverDocuments[T]:
        return self._get_manager("docs", name, model)

    def graph[T: BaseModel](
        self, name: str, model: Type[T] | None = None
    ) -> IBeaverGraph[T]:
        return self._get_manager("graph", name, model)

    def vectors[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverVectors[T]:
        return self._get_manager("vectors", name, model)

    def channel[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverChannel[T]:
        return self._get_manager("channel", name, model)

    def blob[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverBlob[T]:
        return self._get_manager("blob", name, model)

    def log[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverLog[T]:
        return self._get_manager("log", name, model)

    def lock(
        self,
        name: str,
        timeout: float | None = None,
        ttl: float = 60.0,
        poll_interval: float = 0.1,
    ) -> IBeaverLock:
        return self._get_manager("lock", name, timeout, ttl, poll_interval)

    def sketch[T: BaseModel](
        self,
        name: str,
        capacity: int = 1_000_000,
        error_rate: float = 0.01,
        model: type[T] | None = None,
    ) -> IBeaverSketch[T]:
        return self._get_manager("sketch", name, capacity, error_rate, model)

    def events[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverEvents[T]:
        return self._get_manager("events", name, model)
