"""Dataset abstraction with generic schema and configuration support.

Provides generic (engine-agnostic) classes:
- Column: Base column definition
- Schema: Base schema with columns
- SchemaFormat: Display format enum
- DatasetConfig: Configuration for any dataset
- BaseDataset: Abstract interface for all dataset types
- DataRecord: Generic record type
- DataStore: Abstract storage interface

Engine-specific implementations (SQLiteDataset, JSONDataset, etc.) are in their respective modules.
"""

from abc import ABC, abstractmethod
from enum import StrEnum
from typing import Any, Generic, TypeVar

from pydantic import BaseModel, Field


# Type variable for generic data records
T = TypeVar("T")


class SchemaFormat(StrEnum):
    """Format options for schema display."""

    RAW = "raw"  # List of dicts
    PRETTY = "pretty"  # Formatted table


class Column(BaseModel):
    """Generic column definition that works across all engine types.
    
    Base column class for schema definitions. Engine-specific constraints
    and methods should be implemented in subclasses.
    
    Attributes:
        name: Column name
        type: Data type (format depends on engine: SQLite uses TEXT, INTEGER, REAL, etc.)
    """

    name: str = Field(..., description="Column name")
    type: str = Field(..., description="Data type (engine-specific)")


class Schema(BaseModel):
    """Generic table schema with column definitions.
    
    Base schema class that works across all engine types (SQLite, DuckDB, Lance, Qdrant, etc).
    
    Attributes:
        columns: List of Column definitions
    """

    columns: list[Column] = Field(...)


class DataRecord(BaseModel):
    """Generic data record for storage systems.
    
    Attributes:
        id: Record identifier
        data: Record data (can be any JSON-serializable type)
        metadata: Optional metadata dict
    """

    id: str = Field(..., description="Record identifier")
    data: Any = Field(..., description="Record data")
    metadata: dict[str, Any] = Field(default_factory=dict, description="Optional metadata")


class DataStore(ABC, Generic[T]):
    """Abstract base class for data storage implementations.
    
    Type parameter T represents the type of records stored.
    """

    def __init__(self, name: str = "default") -> None:
        """Initialize the data store.
        
        Args:
            name: Name/identifier for this store instance
        """
        self.name = name
        self.record_count = 0

    @abstractmethod
    async def add(self, record: T) -> str:
        """Add a single record to the store.
        
        Args:
            record: The record to add
            
        Returns:
            The record ID
        """
        ...

    @abstractmethod
    async def add_batch(self, records: list[T]) -> list[str]:
        """Add multiple records to the store.
        
        Args:
            records: List of records to add
            
        Returns:
            List of record IDs
        """
        ...

    @abstractmethod
    async def get(self, record_id: str) -> T | None:
        """Retrieve a record by ID.
        
        Args:
            record_id: The ID of the record
            
        Returns:
            The record if found, None otherwise
        """
        ...

    @abstractmethod
    async def query(
        self,
        query_text: str,
        limit: int = 10,
        filters: dict[str, Any] | None = None,
    ) -> list[T]:
        """Query records.
        
        Args:
            query_text: Query text or search string
            limit: Maximum results to return
            filters: Optional additional filters
            
        Returns:
            List of matching records
        """
        ...

    @abstractmethod
    async def delete(self, record_id: str) -> bool:
        """Delete a record from the store.
        
        Args:
            record_id: The ID of the record to delete
            
        Returns:
            True if record was deleted, False if not found
        """
        ...

    @abstractmethod
    async def clear(self) -> None:
        """Delete all records from the store."""
        ...

    @abstractmethod
    async def count(self) -> int:
        """Get total number of records in the store.
        
        Returns:
            Number of records
        """
        ...




class DatasetConfig(BaseModel):
    """Dataset configuration.
    
    Attributes:
        engine: Database engine ("sqlite", "duckdb", "lance", "qdrant", etc.)
        table_name: Name of the table/collection in the database
        in_memory: Whether to use in-memory database
        description: Optional description of the dataset
    """

    engine: str = Field(default="sqlite", description="Database engine")
    table_name: str = Field(..., description="Table name in the database")
    in_memory: bool = Field(default=True, description="Use in-memory storage")
    description: str = Field(default="", description="Dataset description")


class BaseDataset(ABC):
    """Abstract base class for all dataset types.
    
    Implementations must handle engine-specific connection creation
    and data loading. Properties and methods defined here form the
    public dataset interface.
    """

    @property
    @abstractmethod
    def engine(self) -> str:
        """Get the database engine name (sqlite, json, duckdb, lance, qdrant, etc.)."""
        ...

    @property
    @abstractmethod
    def table_name(self) -> str:
        """Get the table/collection name."""
        ...

    @abstractmethod
    def connect(self) -> Any:
        """Get a connection to the dataset.
        
        Returns:
            Connection object specific to the engine type.
            - SQLite: sqlite3.Connection
            - JSON: dict (loaded data)
            - Vector stores: engine-specific client/connection
        """
        ...






