import uuid
import time
from typing import (
    Any,
    Callable,
    Iterator,
    AsyncIterator,
    List,
    Tuple,
    Type,
    Protocol,
    runtime_checkable,
    IO,
    Literal,
    NamedTuple,
)
from pydantic import BaseModel, Field

# --- Data Models ---


class BlobItem(BaseModel):
    key: str
    data: bytes
    metadata: dict | None = None


class ChannelMessage[T](BaseModel):
    channel: str
    payload: T
    timestamp: float


class Document[T](BaseModel):
    id: str = Field(default_factory=lambda: uuid.uuid4().hex)
    body: T


class ScoredDocument[T](BaseModel):
    document: Document[T]
    score: float | None = None


class Event[T](BaseModel):
    id: str = Field(default_factory=lambda: uuid.uuid4().hex)
    event: str
    payload: T
    timestamp: float = Field(default_factory=time.time)


class Edge[T](BaseModel):
    source: str
    target: str
    label: str
    metadata: T | None


class QueueItem[T](NamedTuple):
    priority: float
    timestamp: float
    data: T


class LogEntry[T](NamedTuple):
    timestamp: float
    data: T


class VectorItem[T](BaseModel):
    id: str
    vector: List[float]
    metadata: T | None = None
    score: float = 0


# --- Shared Protocols ---


@runtime_checkable
class IEventHandler(Protocol):
    async def off(self) -> None: ...


@runtime_checkable
class IDocumentQuery[T](Protocol):
    def fts(self, query: str, on: List[str] | None = None) -> "IDocumentQuery[T]": ...
    def fuzzy(self, query: str) -> "IDocumentQuery[T]": ...
    def where(self, *expressions: Any) -> "IDocumentQuery[T]": ...
    def sort(self, **kwargs: Literal["ASC", "DESC"]) -> "IDocumentQuery[T]": ...
    def limit(self, limit: int) -> "IDocumentQuery[T]": ...
    def offset(self, offset: int) -> "IDocumentQuery[T]": ...

    # Async execution
    async def execute(self) -> List[ScoredDocument[T]]: ...
    def __await__(self): ...
    async def __aiter__(self) -> AsyncIterator[ScoredDocument[T]]: ...


@runtime_checkable
class ISketchBatch[T: BaseModel](Protocol):
    def add(self, item: T) -> None: ...
    async def __aenter__(self) -> "ISketchBatch[T]": ...
    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: ...


# --- 1. Blobs ---


@runtime_checkable
class IAsyncBeaverBlob[T: BaseModel](Protocol):
    async def put(
        self, key: str, data: bytes, metadata: dict | None = None
    ) -> None: ...
    async def fetch(self, key: str) -> BlobItem: ...
    async def get(self, key: str) -> bytes: ...
    async def set(self, key: str, data: bytes) -> None: ...
    async def delete(self, key: str) -> None: ...
    async def count(self) -> int: ...
    async def contains(self, key: str) -> bool: ...
    async def clear(self) -> None: ...
    async def keys(self) -> AsyncIterator[str]: ...
    async def items(self) -> AsyncIterator[Tuple[str, bytes]]: ...
    async def dump(
        self, fp: IO[str] | None = None, *, payload: bool = False
    ) -> dict | None: ...
    async def __aiter__(self) -> AsyncIterator[str]: ...


@runtime_checkable
class IBeaverBlob[T: BaseModel](Protocol):
    def put(self, key: str, data: bytes, metadata: dict | None = None) -> None: ...
    def fetch(self, key: str) -> BlobItem: ...
    def get(self, key: str) -> bytes: ...
    def set(self, key: str, data: bytes) -> None: ...
    def delete(self, key: str) -> None: ...
    def count(self) -> int: ...
    def keys(self) -> Iterator[str]: ...
    def items(self) -> Iterator[Tuple[str, bytes]]: ...
    def clear(self) -> None: ...
    def dump(self, fp: IO[str] | None = None) -> dict | None: ...

    def __getitem__(self, key: str) -> bytes: ...
    def __setitem__(self, key: str, data: bytes) -> None: ...
    def __delitem__(self, key: str) -> None: ...
    def __len__(self) -> int: ...
    def __contains__(self, key: str) -> bool: ...
    def __iter__(self) -> Iterator[str]: ...


# --- 2. Channels ---


@runtime_checkable
class IAsyncBeaverChannel[T: BaseModel](Protocol):
    async def publish(self, payload: T) -> None: ...
    def listen(self) -> AsyncIterator[ChannelMessage[T]]: ...
    async def history(self, limit: int = 100) -> list[ChannelMessage[T]]: ...
    async def clear(self) -> None: ...
    async def count(self) -> int: ...


@runtime_checkable
class IBeaverChannel[T: BaseModel](Protocol):
    def publish(self, payload: T) -> None: ...
    def subscribe(self) -> Any: ...  # Returns context manager wrapping Iterator
    def history(self, limit: int = 100) -> list[ChannelMessage[T]]: ...
    def clear(self) -> None: ...
    def count(self) -> int: ...


# --- 3. Dicts ---


@runtime_checkable
class IAsyncBeaverDict[T: BaseModel](Protocol):
    async def set(
        self, key: str, value: T, ttl_seconds: float | None = None
    ) -> None: ...
    async def get(self, key: str) -> T: ...
    async def delete(self, key: str) -> None: ...
    async def fetch(self, key: str, default: Any = None) -> T | Any: ...
    async def pop(self, key: str, default: Any = None) -> T | Any: ...
    async def count(self) -> int: ...
    async def contains(self, key: str) -> bool: ...
    async def clear(self) -> None: ...
    async def keys(self) -> AsyncIterator[str]: ...
    async def values(self) -> AsyncIterator[T]: ...
    async def items(self) -> AsyncIterator[Tuple[str, T]]: ...
    async def dump(self, fp: IO[str] | None = None) -> dict | None: ...
    async def __aiter__(self) -> AsyncIterator[str]: ...


@runtime_checkable
class IBeaverDict[T: BaseModel](Protocol):
    def get(self, key: str) -> T: ...
    def set(self, key: str, value: T, ttl_seconds: float | None = None) -> None: ...
    def delete(self, key: str) -> None: ...
    def fetch(self, key: str, default: Any = None) -> T | Any: ...
    def pop(self, key: str, default: Any = None) -> T | Any: ...
    def keys(self) -> Iterator[str]: ...
    def values(self) -> Iterator[T]: ...
    def items(self) -> Iterator[Tuple[str, T]]: ...
    def clear(self) -> None: ...
    def count(self) -> int: ...
    def dump(self, fp: IO[str] | None = None) -> dict | None: ...

    def __getitem__(self, key: str) -> T: ...
    def __setitem__(self, key: str, value: T) -> None: ...
    def __delitem__(self, key: str) -> None: ...
    def __len__(self) -> int: ...
    def __contains__(self, key: str) -> bool: ...
    def __iter__(self) -> Iterator[str]: ...


# --- 4. Docs ---


@runtime_checkable
class IAsyncBeaverDocuments[T: BaseModel](Protocol):
    async def index(
        self,
        document: Document[T] | None = None,
        id: str | None = None,
        body: T | None = None,
        fts: bool = True,
        fuzzy: bool = False,
    ) -> Document[T]: ...
    async def get(self, id: str) -> Document[T]: ...
    async def get_many(self, ids: List[str]) -> List[Document[T]]: ...
    async def drop(self, id_or_document: str | Document[T]) -> None: ...
    def query(self) -> IDocumentQuery[T]: ...
    async def search(
        self, query: str, on: List[str] | None = None, fuzzy: bool = False
    ) -> List[ScoredDocument[T]]: ...
    async def count(self) -> int: ...
    async def clear(self) -> None: ...
    async def __aiter__(self) -> AsyncIterator[Document[T]]: ...


@runtime_checkable
class IBeaverDocuments[T: BaseModel](Protocol):
    def index(
        self,
        document: Document[T] | None = None,
        id: str | None = None,
        body: T | None = None,
        fts: bool = True,
        fuzzy: bool = False,
    ) -> Document[T]: ...
    def get(self, id: str) -> Document[T]: ...
    def get_many(self, ids: List[str]) -> List[Document[T]]: ...
    def drop(self, id_or_document: str | Document[T]) -> None: ...
    def query(self) -> Any: ...  # Returns sync-wrapped query builder
    def search(
        self, query: str, on: List[str] | None = None, fuzzy: bool = False
    ) -> List[ScoredDocument[T]]: ...
    def count(self) -> int: ...
    def clear(self) -> None: ...
    def __iter__(self) -> Iterator[Document[T]]: ...


# --- 5. Events ---


@runtime_checkable
class IAsyncBeaverEvents[T: BaseModel](Protocol):
    async def attach(self, event: str, callback: Callable) -> IEventHandler: ...
    async def detach(self, event: str, callback: Callable) -> None: ...
    async def emit(self, event: str, payload: T) -> None: ...


@runtime_checkable
class IBeaverEvents[T: BaseModel](Protocol):
    def attach(self, event: str, callback: Callable) -> Any: ...
    def detach(self, event: str, callback: Callable) -> None: ...
    def emit(self, event: str, payload: T) -> None: ...


# --- 6. Graphs ---


@runtime_checkable
class IAsyncBeaverGraph[T: BaseModel](Protocol):
    async def link(
        self, source: str, target: str, label: str, metadata: T | None = None
    ) -> None: ...
    async def unlink(self, source: str, target: str, label: str) -> None: ...
    async def linked(self, source: str, target: str, label: str) -> bool: ...
    async def get(self, source: str, target: str, label: str) -> Edge[T]: ...
    def children(self, source: str, label: str | None = None) -> AsyncIterator[str]: ...
    def parents(self, target: str, label: str | None = None) -> AsyncIterator[str]: ...
    def edges(
        self, source: str, label: str | None = None
    ) -> AsyncIterator[Edge[T]]: ...
    async def count(self) -> int: ...
    async def clear(self) -> None: ...


@runtime_checkable
class IBeaverGraph[T: BaseModel](Protocol):
    def link(
        self, source: str, target: str, label: str, metadata: T | None = None
    ) -> None: ...
    def unlink(self, source: str, target: str, label: str) -> None: ...
    def linked(self, source: str, target: str, label: str) -> bool: ...
    def get(self, source: str, target: str, label: str) -> Edge[T]: ...
    def children(self, source: str, label: str | None = None) -> Iterator[str]: ...
    def parents(self, target: str, label: str | None = None) -> Iterator[str]: ...
    def edges(self, source: str, label: str | None = None) -> Iterator[Edge[T]]: ...
    def count(self) -> int: ...
    def clear(self) -> None: ...


# --- 7. Lists ---


@runtime_checkable
class IAsyncBeaverList[T: BaseModel](Protocol):
    async def get(self, index: int | slice) -> T | List[T]: ...
    async def set(self, index: int, value: T) -> None: ...
    async def delete(self, index: int) -> None: ...
    async def count(self) -> int: ...
    async def contains(self, value: T) -> bool: ...
    async def push(self, value: T) -> None: ...
    async def prepend(self, value: T) -> None: ...
    async def insert(self, index: int, value: T) -> None: ...
    async def pop(self) -> T | None: ...
    async def deque(self) -> T | None: ...
    async def clear(self) -> None: ...
    async def dump(self, fp: IO[str] | None = None) -> dict | None: ...
    async def __aiter__(self) -> AsyncIterator[T]: ...


@runtime_checkable
class IBeaverList[T: BaseModel](Protocol):
    def get(self, index: int | slice) -> T | List[T]: ...
    def set(self, index: int, value: T) -> None: ...
    def delete(self, index: int) -> None: ...
    def count(self) -> int: ...
    def push(self, value: T) -> None: ...
    def prepend(self, value: T) -> None: ...
    def insert(self, index: int, value: T) -> None: ...
    def pop(self) -> T | None: ...
    def deque(self) -> T | None: ...
    def clear(self) -> None: ...
    def dump(self, fp: IO[str] | None = None) -> dict | None: ...

    def __getitem__(self, index: int | slice) -> T | List[T]: ...
    def __setitem__(self, index: int, value: T) -> None: ...
    def __delitem__(self, index: int) -> None: ...
    def __len__(self) -> int: ...
    def __iter__(self) -> Iterator[T]: ...
    def __contains__(self, value: T) -> bool: ...


# --- 8. Locks ---


@runtime_checkable
class IAsyncBeaverLock(Protocol):
    async def acquire(
        self,
        timeout: float | None = None,
        lock_ttl: float | None = None,
        poll_interval: float | None = None,
        block: bool = True,
    ) -> bool: ...
    async def release(self) -> None: ...
    async def renew(self, lock_ttl: float | None = None) -> bool: ...
    async def clear(self) -> bool: ...
    async def __aenter__(self) -> "IAsyncBeaverLock": ...
    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: ...


@runtime_checkable
class IBeaverLock(Protocol):
    def acquire(
        self,
        timeout: float | None = None,
        lock_ttl: float | None = None,
        poll_interval: float | None = None,
        block: bool = True,
    ) -> bool: ...
    def release(self) -> None: ...
    def renew(self, lock_ttl: float | None = None) -> bool: ...
    def clear(self) -> bool: ...
    def __enter__(self) -> "IBeaverLock": ...
    def __exit__(self, exc_type, exc_val, exc_tb) -> None: ...


# --- 9. Logs ---


@runtime_checkable
class IAsyncBeaverLog[T: BaseModel](Protocol):
    async def log(self, data: T, timestamp: float | None = None) -> None: ...
    async def range(
        self,
        start: float | None = None,
        end: float | None = None,
        limit: int | None = None,
    ) -> list[LogEntry[T]]: ...
    def live(self, poll_interval: float = 0.1) -> AsyncIterator[LogEntry[T]]: ...
    async def count(self) -> int: ...
    async def clear(self) -> None: ...
    async def dump(self, fp: IO[str] | None = None) -> dict | None: ...


@runtime_checkable
class IBeaverLog[T: BaseModel](Protocol):
    def log(self, data: T, timestamp: float | None = None) -> None: ...
    def range(
        self,
        start: float | None = None,
        end: float | None = None,
        limit: int | None = None,
    ) -> list[LogEntry[T]]: ...
    def live(self, poll_interval: float = 0.1) -> Iterator[LogEntry[T]]: ...
    def count(self) -> int: ...
    def clear(self) -> None: ...
    def dump(self, fp: IO[str] | None = None) -> dict | None: ...


# --- 10. Queues ---


@runtime_checkable
class IAsyncBeaverQueue[T: BaseModel](Protocol):
    async def put(self, data: T, priority: float) -> None: ...
    async def peek(self) -> QueueItem[T] | None: ...
    async def get(
        self, block: bool = True, timeout: float | None = None
    ) -> QueueItem[T]: ...
    async def count(self) -> int: ...
    async def clear(self) -> None: ...
    async def dump(self, fp: IO[str] | None = None) -> dict | None: ...
    async def __aiter__(self) -> AsyncIterator[QueueItem[T]]: ...


@runtime_checkable
class IBeaverQueue[T: BaseModel](Protocol):
    def put(self, data: T, priority: float) -> None: ...
    def peek(self) -> QueueItem[T] | None: ...
    def get(self, block: bool = True, timeout: float | None = None) -> QueueItem[T]: ...
    def count(self) -> int: ...
    def clear(self) -> None: ...
    def dump(self, fp: IO[str] | None = None) -> dict | None: ...
    def __len__(self) -> int: ...
    def __iter__(self) -> Iterator[QueueItem[T]]: ...


# --- 11. Sketches ---


@runtime_checkable
class IAsyncBeaverSketch[T: BaseModel](Protocol):
    async def add(self, item: T) -> None: ...
    async def contains(self, item: T) -> bool: ...
    async def count(self) -> int: ...
    def batched(self) -> ISketchBatch: ...
    async def clear(self) -> None: ...


@runtime_checkable
class IBeaverSketch[T: BaseModel](Protocol):
    def add(self, item: T) -> None: ...
    def contains(self, item: T) -> bool: ...
    def count(self) -> int: ...
    def batched(self) -> Any: ...  # Sync Context Manager
    def clear(self) -> None: ...
    def __len__(self) -> int: ...
    def __contains__(self, item: T) -> bool: ...


# --- 12. Vectors ---


@runtime_checkable
class IAsyncBeaverVectors[T: BaseModel](Protocol):
    async def set(
        self, id: str, vector: List[float], metadata: T | None = None
    ) -> None: ...
    async def get(self, id: str) -> VectorItem[T]: ...
    async def delete(self, id: str) -> None: ...
    async def near(
        self,
        vector: List[float],
        k: int = 10,
        candidate_ids: List[str] | None = None,
        filters: List[Any] | None = None,
        metric: Any = None,
    ) -> List[VectorItem[T]]: ...
    async def count(self) -> int: ...
    async def clear(self) -> None: ...
    async def __aiter__(self) -> AsyncIterator[VectorItem[T]]: ...


@runtime_checkable
class IBeaverVectors[T: BaseModel](Protocol):
    def set(self, id: str, vector: List[float], metadata: T | None = None) -> None: ...
    def get(self, id: str) -> VectorItem[T]: ...
    def delete(self, id: str) -> None: ...
    def near(
        self,
        vector: List[float],
        k: int = 10,
        candidate_ids: List[str] | None = None,
        filters: List[Any] | None = None,
        metric: Any = None,
    ) -> List[VectorItem[T]]: ...
    def count(self) -> int: ...
    def clear(self) -> None: ...
    def __iter__(self) -> Iterator[VectorItem[T]]: ...


# --- Main Database Factories ---


@runtime_checkable
class IAsyncBeaverDB(Protocol):
    async def connect(self) -> "IAsyncBeaverDB": ...
    async def close(self) -> None: ...

    def dict[T: BaseModel](
        self, name: str, model: type[T] | None = None, secret: str | None = None
    ) -> IAsyncBeaverDict[T]: ...
    def list[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IAsyncBeaverList[T]: ...
    def queue[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IAsyncBeaverQueue[T]: ...
    def docs[T: BaseModel](
        self, name: str, model: Type[T] | None = None
    ) -> IAsyncBeaverDocuments[T]: ...
    def graph[T: BaseModel](
        self, name: str, model: Type[T] | None = None
    ) -> IAsyncBeaverGraph[T]: ...
    def vectors[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IAsyncBeaverVectors[T]: ...
    def channel[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IAsyncBeaverChannel[T]: ...
    def blob[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IAsyncBeaverBlob[T]: ...
    def log[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IAsyncBeaverLog[T]: ...
    def lock(
        self,
        name: str,
        timeout: float | None = None,
        ttl: float = 60.0,
        poll_interval: float = 0.1,
    ) -> IAsyncBeaverLock: ...
    def sketch[T: BaseModel](
        self,
        name: str,
        capacity: int = 1_000_000,
        error_rate: float = 0.01,
        model: type[T] | None = None,
    ) -> IAsyncBeaverSketch[T]: ...
    def events[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IAsyncBeaverEvents[T]: ...


@runtime_checkable
class IBeaverDB(Protocol):
    def close(self) -> None: ...
    def __enter__(self) -> "IBeaverDB": ...
    def __exit__(self, exc_type, exc_val, exc_tb) -> None: ...

    def dict[T: BaseModel](
        self, name: str, model: type[T] | None = None, secret: str | None = None
    ) -> IBeaverDict[T]: ...
    def list[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverList[T]: ...
    def queue[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverQueue[T]: ...
    def docs[T: BaseModel](
        self, name: str, model: Type[T] | None = None
    ) -> IBeaverDocuments[T]: ...
    def graph[T: BaseModel](
        self, name: str, model: Type[T] | None = None
    ) -> IBeaverGraph[T]: ...
    def vectors[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverVectors[T]: ...
    def channel[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverChannel[T]: ...
    def blob[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverBlob[T]: ...
    def log[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverLog[T]: ...
    def lock(
        self,
        name: str,
        timeout: float | None = None,
        ttl: float = 60.0,
        poll_interval: float = 0.1,
    ) -> IBeaverLock: ...
    def sketch[T: BaseModel](
        self,
        name: str,
        capacity: int = 1_000_000,
        error_rate: float = 0.01,
        model: type[T] | None = None,
    ) -> IBeaverSketch[T]: ...
    def events[T: BaseModel](
        self, name: str, model: type[T] | None = None
    ) -> IBeaverEvents[T]: ...
