"""Type stubs for aifastdb - IDE autocomplete support."""

from typing import Any, overload

import numpy as np
import numpy.typing as npt

# ---- Version ----
__version__: str

# ---- Exception Hierarchy ----

class AiFastDbError(RuntimeError):
    """Base exception for all AiFastDb errors."""
    ...

class StoreNotFoundError(AiFastDbError):
    """Raised when a store is not found by alias."""
    ...

class RecordNotFoundError(AiFastDbError):
    """Raised when a record is not found by ID."""
    ...

class QueryError(AiFastDbError):
    """Raised on query execution failure."""
    ...

class WriteError(AiFastDbError):
    """Raised on write operation failure."""
    ...

class TransactionError(AiFastDbError):
    """Raised on transaction failure."""
    ...

class StorageError(AiFastDbError):
    """Raised on storage/persistence failure."""
    ...

class CircuitBreakerError(AiFastDbError):
    """Raised when a circuit breaker is open for a store."""
    ...

class RateLimitError(AiFastDbError):
    """Raised when rate limit is exceeded."""
    ...

# ---- Enums ----

class ContentType:
    """Content type for stored records."""
    Text: int
    Json: int
    Binary: int
    Embedding: int
    Code: int

class StoreType:
    """Store type in the federation."""
    Graph: int
    Document: int
    Vector: int

# ---- Database classes (phase-24) ----

class Collection:
    """A single named collection in the AiFastDb database.

    Obtained via ``db.collection("name")``.
    """

    @property
    def name(self) -> str: ...
    def document_count(self) -> int: ...
    def vector_count(self) -> int: ...

    def remember(
        self,
        id: str,
        content: str,
        *,
        content_type: str = "text",
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        embedding: list[float] | None = None,
        importance: float = 1.0,
        parent_id: str | None = None,
    ) -> None: ...
    def get(self, id: str) -> dict[str, Any] | None: ...
    def forget(self, id: str) -> bool: ...
    def touch(self, id: str) -> bool: ...
    def recall(
        self,
        *,
        query: str | None = None,
        embedding: list[float] | None = None,
        limit: int = 10,
        min_importance: float | None = None,
    ) -> list[dict[str, Any]]: ...
    def recall_rrf(
        self,
        *,
        query: str | None = None,
        embedding: list[float] | None = None,
        limit: int = 10,
        min_importance: float | None = None,
        k: float = 60.0,
    ) -> list[dict[str, Any]]: ...
    def recall_all(self) -> list[dict[str, Any]]: ...
    def flush(self) -> None: ...
    def __repr__(self) -> str: ...
    def __len__(self) -> int: ...


class AiFastDb:
    """High-level Python interface to the AiFastDb database.

    Supports collection management, CRUD, hybrid search, and batch operations.

    Example:
        >>> with AiFastDb("/path/to/db") as db:
        ...     db.remember("id1", "hello world", tags=["greeting"])
        ...     result = db.get("id1")
    """

    def __init__(self, path: str, *, dimension: int = 128) -> None: ...

    def __enter__(self) -> "AiFastDb": ...
    def __exit__(self, *args: Any) -> bool: ...

    @property
    def path(self) -> str: ...
    @property
    def dimension(self) -> int: ...

    # Collection management
    def list_collections(self) -> list[str]: ...
    def has_collection(self, name: str) -> bool: ...
    def delete_collection(self, name: str) -> bool: ...
    def collection(
        self,
        name: str,
        *,
        dimension: int | None = None,
        description: str | None = None,
    ) -> Collection: ...

    # CRUD (on default collection)
    def remember(
        self,
        id: str,
        content: str,
        *,
        content_type: str = "text",
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        embedding: list[float] | None = None,
        importance: float = 1.0,
        parent_id: str | None = None,
    ) -> None: ...
    def get(self, id: str) -> dict[str, Any] | None: ...
    def forget(self, id: str) -> bool: ...
    def touch(self, id: str) -> bool: ...

    # Batch operations
    def remember_many(self, items: list[dict[str, Any]]) -> int: ...
    def get_many(self, ids: list[str]) -> dict[str, dict[str, Any]]: ...

    # Search
    def recall(
        self,
        *,
        query: str | None = None,
        embedding: list[float] | None = None,
        limit: int = 10,
        min_importance: float | None = None,
    ) -> list[dict[str, Any]]: ...
    def recall_rrf(
        self,
        *,
        query: str | None = None,
        embedding: list[float] | None = None,
        limit: int = 10,
        min_importance: float | None = None,
        k: float = 60.0,
    ) -> list[dict[str, Any]]: ...
    def recall_all(self) -> list[dict[str, Any]]: ...

    # Utility
    def flush(self) -> None: ...
    def document_count(self) -> int: ...
    def __repr__(self) -> str: ...
    def __len__(self) -> int: ...

# phase-25: SocialGraphV2
class SocialGraphV2:
    """Social graph engine for managing persons, relationships, tags, and entities."""

    def __init__(self, path: str, *, shard_count: int = 0, wal_enabled: bool = True) -> None: ...

    # T25.1: Lifecycle
    def recover(self) -> None: ...
    def shutdown(self) -> None: ...
    def __enter__(self) -> "SocialGraphV2": ...
    def __exit__(self, *args: Any) -> None: ...

    @property
    def path(self) -> str: ...

    # T25.2: Person CRUD
    def add_person(self, input: dict[str, Any]) -> dict[str, Any]: ...
    def get_person(self, id: str) -> dict[str, Any] | None: ...
    def update_person(self, id: str, updates: dict[str, Any]) -> dict[str, Any] | None: ...
    def delete_person(self, id: str) -> bool: ...
    def list_persons(self, filter: dict[str, Any] | None = None) -> list[dict[str, Any]]: ...

    # T25.3: Friend / Relation Management
    def add_friend(self, person_a: str, person_b: str, config: dict[str, Any] | None = None) -> str | None: ...
    def remove_friend(self, person_a: str, person_b: str) -> bool: ...
    def get_friends(self, person_id: str) -> list[dict[str, Any]]: ...
    def are_friends(self, person_a: str, person_b: str) -> bool: ...
    def get_direct_relation(self, person_a: str, person_b: str) -> dict[str, Any] | None: ...

    # T25.4: Tag Management
    def add_tag(self, input: dict[str, Any]) -> dict[str, Any]: ...
    def get_tag(self, id: str) -> dict[str, Any] | None: ...
    def tag_person(self, person_id: str, tag_id: str) -> str | None: ...
    def untag_person(self, person_id: str, tag_id: str) -> bool: ...
    def get_person_tags(self, person_id: str) -> list[dict[str, Any]]: ...
    def list_tags(self, filter: dict[str, Any] | None = None) -> list[dict[str, Any]]: ...

    # T25.5: Graph Export
    def export_graph(self, options: dict[str, Any] | None = None) -> dict[str, Any]: ...
    def get_graph_stats(self) -> dict[str, Any]: ...

    # T25.6: Generic Entity API
    def add_entity(self, name: str, entity_type: str, properties: dict[str, Any] | None = None) -> dict[str, Any]: ...
    def update_entity(self, id: str, updates: dict[str, Any]) -> dict[str, Any] | None: ...
    def delete_entity(self, id: str) -> bool: ...
    def list_entities_by_type(self, entity_type: str) -> list[dict[str, Any]]: ...
    def get_entity_types(self) -> list[str]: ...
    def contains_entity(self, id: str) -> bool: ...
    def entity_count(self) -> int: ...
    def relation_count(self) -> int: ...

    # T25.7: WAL & Metrics
    def flush(self) -> None: ...
    def sync(self) -> None: ...
    def is_wal_enabled(self) -> bool: ...
    def get_wal_diagnostics(self) -> dict[str, Any]: ...
    def get_metrics(self) -> dict[str, Any]: ...
    def shard_count(self) -> int: ...

    # T26.1: Path Finding
    def find_path(self, from_id: str, to_id: str, options: dict[str, Any] | None = None) -> dict[str, Any] | None: ...
    def find_all_paths(self, from_id: str, to_id: str, max_hops: int = 6, max_paths: int = 10) -> dict[str, Any]: ...
    def find_intermediaries(self, person_a: str, person_b: str, max_hops: int = 6) -> list[dict[str, Any]]: ...

    # T26.2: Network Analysis
    def get_mutual_friends(self, person_a: str, person_b: str) -> list[dict[str, Any]]: ...
    def calculate_relationship_strength(self, person_a: str, person_b: str) -> dict[str, Any]: ...
    def suggest_connections(self, person_id: str, config: dict[str, Any] | None = None) -> list[dict[str, Any]]: ...

    # T26.3: N-Degree Contacts & Network Stats
    def get_nth_degree_contacts(self, person_id: str, max_degree: int = 3) -> dict[str, Any]: ...
    def get_network_stats(self, person_id: str) -> dict[str, Any]: ...

    # T26.4: Item Management
    def add_item(self, input: dict[str, Any]) -> dict[str, Any]: ...
    def add_item_to_person(self, person_id: str, item_id: str, relation_type: str | None = None) -> str | None: ...
    def remove_item_from_person(self, person_id: str, item_id: str) -> bool: ...
    def get_person_items(self, person_id: str) -> list[dict[str, Any]]: ...
    def find_shared_items(self, person_a: str, person_b: str) -> list[dict[str, Any]]: ...
    def find_people_by_shared_items(self, person_id: str, category: str | None = None) -> list[dict[str, Any]]: ...
    def list_items(self, category: str | None = None) -> list[dict[str, Any]]: ...

    # T26.5: Attribute Management
    def add_attribute(self, input: dict[str, Any]) -> dict[str, Any]: ...
    def find_or_create_attribute(self, category: str, value: str) -> dict[str, Any]: ...
    def set_person_attribute(self, person_id: str, attr_id: str, relation_type: str | None = None) -> str | None: ...
    def remove_person_attribute(self, person_id: str, attr_id: str) -> bool: ...
    def get_person_attributes(self, person_id: str) -> list[dict[str, Any]]: ...
    def list_attributes(self, category: str | None = None) -> list[dict[str, Any]]: ...
    def find_alumni(self, person_id: str) -> list[dict[str, Any]]: ...
    def find_same_city(self, person_id: str) -> list[dict[str, Any]]: ...
    def find_shared_attributes(self, person_a: str, person_b: str) -> list[dict[str, Any]]: ...

    # T26.6: Birthday & Family Tree
    def set_birthday(self, person_id: str, birthday: str) -> bool: ...
    def get_birthday(self, person_id: str) -> str | None: ...
    def find_upcoming_birthdays(self, days: int = 30) -> list[dict[str, Any]]: ...
    def add_parent_child(self, parent_id: str, child_id: str) -> str | None: ...
    def get_ancestors(self, person_id: str, max_generations: int = 5) -> list[dict[str, Any]]: ...
    def get_descendants(self, person_id: str, max_generations: int = 5) -> list[dict[str, Any]]: ...
    def get_parents(self, person_id: str) -> list[dict[str, Any]]: ...
    def get_children(self, person_id: str) -> list[dict[str, Any]]: ...
    def add_spouse(self, person_a: str, person_b: str) -> str | None: ...
    def add_sibling(self, person_a: str, person_b: str, half_sibling: bool = False) -> str | None: ...

    # T26.7: Compatibility Analysis
    def analyze_compatibility(self, person_a: str, person_b: str) -> dict[str, Any]: ...

# phase-27: DocumentStore & Federation

class DocumentStore:
    """Enhanced JSONL document store.

    Usage:
        with DocumentStore("/path/to/data.jsonl") as store:
            store.put("doc1", "Hello", doc_type="text", tags=["greeting"])
            doc = store.get("doc1")
    """

    def __init__(
        self,
        path: str,
        *,
        flush_interval_ms: int = 1000,
        enable_indexing: bool = True,
        enable_file_lock: bool = True,
    ) -> None: ...

    # Context manager
    def __enter__(self) -> "DocumentStore": ...
    def __exit__(self, *args: Any) -> bool: ...

    @property
    def path(self) -> str: ...
    def __repr__(self) -> str: ...
    def __len__(self) -> int: ...

    # CRUD
    def put(
        self,
        id: str,
        content: str,
        *,
        doc_type: str = "text",
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        importance: float = 0.5,
    ) -> None: ...
    def get(self, id: str) -> dict[str, Any] | None: ...
    def has(self, id: str) -> bool: ...
    def delete(self, id: str) -> dict[str, Any] | None: ...

    # Query
    def find_by_tag(self, tag: str) -> list[dict[str, Any]]: ...
    def find_by_type(self, doc_type: str) -> list[dict[str, Any]]: ...
    def find_by_time_range(self, start: int, end: int) -> list[dict[str, Any]]: ...
    def find_by_min_importance(self, min_importance: float) -> list[dict[str, Any]]: ...

    # Utilities
    def ids(self) -> list[str]: ...
    def is_empty(self) -> bool: ...
    def all_tags(self) -> list[str]: ...
    def all(self) -> list[dict[str, Any]]: ...
    def flush(self) -> None: ...
    def compact(self) -> None: ...


class FederatedTransaction:
    """Cross-store transaction with 2PC.

    Usage:
        tx = federation.begin_transaction(["social", "docs"])
        tx.write("social", {"op": "CreateEntity", ...})
        tx.write("docs", {"op": "CreateDocument", ...})
        tx.commit()
    """

    def write(self, store_alias: str, operation: dict[str, Any]) -> None: ...
    def commit(self) -> None: ...
    def rollback(self) -> None: ...


class Federation:
    """Heterogeneous storage federation — unified interface for graph, document, and vector stores.

    Usage:
        with Federation() as fed:
            fed.register_graph("social", "/path/to/social")
            fed.register_document("docs", "/path/to/docs.jsonl")
            result = fed.query("social", {"target": "person"})
    """

    def __init__(
        self,
        *,
        enable_permissions: bool = False,
        query_timeout_ms: int = 30000,
    ) -> None: ...

    # Context manager
    def __enter__(self) -> "Federation": ...
    def __exit__(self, *args: Any) -> bool: ...
    def __repr__(self) -> str: ...

    # T27.3: Store registration
    def register_graph(
        self,
        alias: str,
        path: str,
        *,
        shard_count: int = 0,
        wal_enabled: bool = True,
    ) -> None: ...
    def register_document(self, alias: str, path: str) -> None: ...
    def register_vector(self, alias: str, path: str) -> None: ...
    def unregister(self, alias: str) -> None: ...
    def list_stores(self) -> list[tuple[str, str]]: ...

    # T27.4: Unified query
    def query(self, store_alias: str, query: dict[str, Any]) -> dict[str, Any]: ...
    def cross_query(self, query: dict[str, Any]) -> dict[str, Any]: ...

    # T27.5: Unified write
    def write(self, store_alias: str, operation: dict[str, Any]) -> str: ...
    def batch_write(self, store_alias: str, operations: list[dict[str, Any]]) -> dict[str, Any]: ...

    # T27.6: Transaction support
    def begin_transaction(self, store_aliases: list[str]) -> FederatedTransaction: ...

    # T27.7: Permission-checked operations
    def write_with_auth(
        self,
        store_alias: str,
        operation: dict[str, Any],
        caller: str,
        roles: list[str],
    ) -> str: ...
    def query_with_auth(
        self,
        store_alias: str,
        query: dict[str, Any],
        caller: str,
        roles: list[str],
    ) -> dict[str, Any]: ...
    def get_record_actions(
        self,
        store_alias: str,
        record_id: str,
        caller: str,
        roles: list[str],
    ) -> dict[str, Any]: ...

    # Lifecycle
    def sync_all(self) -> None: ...
    def close(self) -> None: ...
    def rebuild_index(self) -> None: ...


# ---- Phase-28: Perception / Embedding / Reasoning ----

class EmbeddingEngine:
    """Embedding engine for generating vector representations.

    Wraps the Rust PerceptionEngine trait. Supports text, batch,
    and multimodal embedding with optional numpy output.

    Example:
        >>> engine = EmbeddingEngine()           # dummy, 384d
        >>> vec = engine.embed("Hello world")
        >>> arr = engine.embed_numpy("Hello")    # numpy.ndarray
    """

    def __init__(
        self,
        *,
        engine_type: str = "dummy",
        model_id: str | None = None,
        dimension: int | None = None,
        device: str = "cpu",
        cache_dir: str | None = None,
    ) -> None: ...

    # Preset constructors
    @staticmethod
    def dummy(dimension: int = 384) -> "EmbeddingEngine": ...
    @staticmethod
    def minilm() -> "EmbeddingEngine": ...
    @staticmethod
    def bge_small() -> "EmbeddingEngine": ...
    @staticmethod
    def gemma(dimension: int = 768) -> "EmbeddingEngine": ...

    # T28.1: Embed
    def embed(self, text: str) -> list[float]: ...
    def embed_batch(self, texts: list[str]) -> list[list[float]]: ...

    # T28.4: numpy
    def embed_numpy(self, text: str) -> npt.NDArray[np.float32]: ...
    def embed_batch_numpy(self, texts: list[str]) -> list[npt.NDArray[np.float32]]: ...

    # T28.2: Multimodal
    def embed_multimodal(self, data: bytes, mime_type: str) -> list[float]: ...
    def embed_multimodal_numpy(self, data: bytes, mime_type: str) -> npt.NDArray[np.float32]: ...

    # Utility
    @staticmethod
    def cosine_similarity(a: list[float], b: list[float]) -> float: ...

    # Properties
    @property
    def dimension(self) -> int: ...
    @property
    def model_name(self) -> str: ...
    @property
    def max_input_length(self) -> int | None: ...
    @property
    def supported_modalities(self) -> list[str]: ...

    def __repr__(self) -> str: ...


class ReasoningEngine:
    """Reasoning engine for LLM-based cognitive tasks.

    Connects to external LLM APIs (Ollama, OpenAI, etc.) for
    intent understanding, logic checking, and summarization.

    Example:
        >>> engine = ReasoningEngine.ollama("llama3.2")
        >>> result = engine.understand_intent("What's the weather?")
    """

    def __init__(
        self,
        provider: str,
        model: str,
        *,
        endpoint: str | None = None,
        api_key: str | None = None,
        timeout: int = 60,
        max_tokens: int | None = None,
        temperature: float | None = None,
    ) -> None: ...

    # Preset constructors
    @staticmethod
    def ollama(model: str = "llama3.2", *, endpoint: str | None = None, timeout: int = 60) -> "ReasoningEngine": ...
    @staticmethod
    def openai(api_key: str, model: str = "gpt-4o-mini", *, timeout: int = 60, max_tokens: int | None = None) -> "ReasoningEngine": ...

    # T28.5: Reasoning methods
    def check_logic(self, fact_a: str, fact_b: str) -> bool: ...
    def understand_intent(self, query: str) -> str: ...
    def summarize(self, content: str, *, max_tokens: int = 200) -> str: ...
    def chat(self, messages: list[dict[str, str]]) -> dict[str, Any]: ...

    @property
    def model_name(self) -> str: ...
    @property
    def is_available(self) -> bool: ...

    def __repr__(self) -> str: ...


class PerceptionPipeline:
    """Combined perception (embedding) + reasoning pipeline.

    Provides a unified interface for embedding and reasoning workflows,
    including the convenience method ``remember_with_embedding()``.

    Example:
        >>> pipe = PerceptionPipeline()
        >>> result = pipe.remember_with_embedding("Hello world", "doc1")
        >>> print(result["embedding"][:5])
    """

    def __init__(
        self,
        *,
        engine_type: str = "dummy",
        dimension: int = 384,
        device: str = "cpu",
        reasoning_provider: str | None = None,
        reasoning_model: str | None = None,
        reasoning_endpoint: str | None = None,
        reasoning_api_key: str | None = None,
    ) -> None: ...

    # Embedding
    def embed(self, text: str) -> list[float]: ...
    def embed_batch(self, texts: list[str]) -> list[list[float]]: ...
    def embed_numpy(self, text: str) -> npt.NDArray[np.float32]: ...
    def embed_multimodal(self, data: bytes, mime_type: str) -> list[float]: ...

    # T28.3: Auto-embed and remember
    def remember_with_embedding(
        self,
        content: str,
        id: str,
        *,
        tags: list[str] | None = None,
        doc_type: str = "text",
    ) -> dict[str, Any]: ...

    def remember_with_embedding_numpy(
        self,
        content: str,
        id: str,
        *,
        tags: list[str] | None = None,
        doc_type: str = "text",
    ) -> dict[str, Any]: ...

    # Reasoning (requires reasoning engine)
    def check_logic(self, fact_a: str, fact_b: str) -> bool: ...
    def understand_intent(self, query: str) -> str: ...
    def summarize(self, content: str, *, max_tokens: int = 200) -> str: ...

    # Properties
    @property
    def dimension(self) -> int: ...
    @property
    def has_reasoning(self) -> bool: ...
    @property
    def model_name(self) -> str: ...

    def __repr__(self) -> str: ...

# =============================================================================
# High-concurrency storage (phase-30)
# =============================================================================

class ConcurrentDocumentStore:
    """High-concurrency document store with sharded write queues, WAL, and backpressure."""

    def __init__(
        self,
        path: str,
        *,
        shard_count: int = 0,
        enable_wal: bool = True,
        mode: str = "default",
    ) -> None: ...

    # Write
    def write(
        self,
        id: str,
        content: str,
        *,
        tags: list[str] = ...,
        doc_type: str = "text",
        importance: float = 0.5,
    ) -> None: ...
    def write_sync(
        self,
        id: str,
        content: str,
        *,
        tags: list[str] = ...,
        doc_type: str = "text",
        importance: float = 0.5,
    ) -> None: ...
    def write_batch(self, docs: list[dict]) -> int: ...

    # Read
    def get(self, id: str) -> dict | None: ...
    def has(self, id: str) -> bool: ...
    def __len__(self) -> int: ...
    def is_empty(self) -> bool: ...
    def ids(self) -> list[str]: ...

    # Management
    def flush(self) -> None: ...
    def metrics(self) -> dict: ...
    def queue_fill_ratio(self) -> float: ...
    def recover(self) -> dict: ...
    @property
    def shard_count(self) -> int: ...
    def close(self) -> None: ...

    # Context manager
    def __enter__(self) -> ConcurrentDocumentStore: ...
    def __exit__(self, exc_type: type | None = None, exc_val: BaseException | None = None, exc_tb: object | None = None) -> bool: ...
    def __repr__(self) -> str: ...


class VectorIndex:
    """High-concurrency sharded HNSW vector index."""

    def __init__(
        self,
        dimension: int,
        *,
        shard_count: int = 0,
        mode: str = "default",
        m: int = 16,
        ef_construction: int = 100,
        ef_search: int = 50,
        max_elements: int = 100000,
    ) -> None: ...

    # Write
    def insert(self, id: str, vector: list[float]) -> None: ...
    def insert_sync(self, id: str, vector: list[float]) -> None: ...
    def insert_batch(self, items: list[tuple[str, list[float]]]) -> int: ...
    def delete(self, id: str) -> None: ...

    # Search
    def search(self, query: list[float], k: int = 10) -> list[dict]: ...
    def search_with_ef(self, query: list[float], k: int = 10, ef: int = 200) -> list[dict]: ...
    def search_numpy(self, query: numpy.typing.NDArray[numpy.float32], k: int = 10) -> list[dict]: ...
    def get(self, id: str) -> list[float] | None: ...
    def get_numpy(self, id: str) -> numpy.typing.NDArray[numpy.float32] | None: ...
    def has(self, id: str) -> bool: ...
    def __len__(self) -> int: ...
    def is_empty(self) -> bool: ...

    # Management
    def flush(self) -> None: ...
    def metrics(self) -> dict: ...
    def queue_fill_ratio(self) -> float: ...
    @property
    def shard_count(self) -> int: ...
    def shard_stats(self) -> list[tuple[int, int]]: ...
    @property
    def dimension(self) -> int: ...
    def backpressure_status(self) -> str: ...
    def close(self) -> None: ...

    # Context manager
    def __enter__(self) -> VectorIndex: ...
    def __exit__(self, exc_type: type | None = None, exc_val: BaseException | None = None, exc_tb: object | None = None) -> bool: ...
    def __repr__(self) -> str: ...
