"""Pit - A searchable collection of objects backed by DuckDB."""

from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, TypeVar

import duckdb
from pydantic import BaseModel

from mosh.metadata import ObjectMetadata, StoredObject

T = TypeVar("T", bound=BaseModel)


@dataclass
class DuckDBConfig:
    """Configuration for DuckDB performance tuning.

    DuckDB uses RAM for fast query processing. These settings control
    memory usage and parallelism for optimal performance.

    Attributes:
        memory_limit: Maximum memory for DuckDB operations (e.g., "4GB", "512MB").
            Default None uses DuckDB's default (80% of available RAM).
        threads: Number of threads for parallel query execution.
            Default None uses DuckDB's default (number of CPU cores).
        temp_directory: Directory for temporary files when memory is exceeded.
            Default None uses system temp directory.
        preserve_insertion_order: Whether to preserve row insertion order.
            Default True for predictable iteration.
    """

    memory_limit: str | None = None
    threads: int | None = None
    temp_directory: str | None = None
    preserve_insertion_order: bool = True


class Pit:
    """A searchable collection of pydantic BaseModel objects.

    Uses DuckDB as the storage backend. Each pit is stored in a single DuckDB database
    with objects serialized as JSON. Metadata (id, timestamps, tags) is managed
    separately from the user's object data.

    DuckDB leverages local RAM for high-performance query processing. Use DuckDBConfig
    to tune memory limits, parallelism, and other performance settings.
    """

    def __init__(
        self,
        name: str,
        db_path: str | None = None,
        object_types: list[type[BaseModel]] | None = None,
        id_field: str | None = None,
        duckdb_config: DuckDBConfig | None = None,
    ):
        """Initialize pit with DuckDB backend.

        Args:
            name: Name of the pit (used as table name).
            db_path: Path to DuckDB database file. If None, uses in-memory database.
            object_types: List of BaseModel subclasses this pit will store.
                         Used for deserializing objects.
            id_field: Optional field name to extract ID from objects. If specified,
                     the value of this field will be used as the object ID instead
                     of auto-generating a UUID.
            duckdb_config: Optional DuckDB performance configuration. Controls memory
                          limits, thread count, and other performance settings.
        """
        self.name = name
        self.db_path = db_path
        self.id_field = id_field
        self.duckdb_config = duckdb_config or DuckDBConfig()
        self._object_types: dict[str, type[BaseModel]] = {}

        if object_types:
            for obj_type in object_types:
                self._object_types[obj_type.__name__] = obj_type

        # Connect to DuckDB
        if db_path:
            Path(db_path).parent.mkdir(parents=True, exist_ok=True)
            self._conn = duckdb.connect(db_path)
        else:
            self._conn = duckdb.connect(":memory:")

        # Apply DuckDB configuration for RAM-based performance
        self._apply_duckdb_config()

        self._init_schema()

    def _apply_duckdb_config(self) -> None:
        """Apply DuckDB configuration settings for performance tuning."""
        config = self.duckdb_config

        if config.memory_limit:
            self._conn.execute(f"SET memory_limit = '{config.memory_limit}'")

        if config.threads is not None:
            self._conn.execute(f"SET threads = {config.threads}")

        if config.temp_directory:
            self._conn.execute(f"SET temp_directory = '{config.temp_directory}'")

        if config.preserve_insertion_order:
            self._conn.execute("SET preserve_insertion_order = true")
        else:
            self._conn.execute("SET preserve_insertion_order = false")

    def _init_schema(self) -> None:
        """Initialize the database schema."""
        self._conn.execute(f"""
            CREATE TABLE IF NOT EXISTS {self.name} (
                id VARCHAR PRIMARY KEY,
                type_name VARCHAR NOT NULL,
                created_at TIMESTAMP,
                updated_at TIMESTAMP,
                tags VARCHAR[],
                data JSON NOT NULL
            )
        """)
        # Create indexes for common queries
        self._conn.execute(f"CREATE INDEX IF NOT EXISTS idx_{self.name}_type ON {self.name}(type_name)")
        self._conn.execute(f"CREATE INDEX IF NOT EXISTS idx_{self.name}_created ON {self.name}(created_at)")
        self._conn.execute(f"CREATE INDEX IF NOT EXISTS idx_{self.name}_updated ON {self.name}(updated_at)")

    def register_type(self, obj_type: type[BaseModel]) -> None:
        """Register an object type for deserialization."""
        self._object_types[obj_type.__name__] = obj_type

    def _get_id_for_object(self, obj: BaseModel, explicit_id: str | None) -> str:
        """Determine the ID for an object.

        Priority:
        1. Explicit ID passed to put()
        2. Value from id_field if configured
        3. Auto-generated UUID
        """
        if explicit_id is not None:
            return explicit_id
        if self.id_field is not None:
            value = getattr(obj, self.id_field, None)
            if value is not None:
                return str(value)
        # Generate new UUID
        from uuid import uuid4

        return str(uuid4())

    def _serialize(self, obj: BaseModel, metadata: ObjectMetadata) -> dict[str, Any]:
        """Serialize a BaseModel object for storage."""
        return {
            "id": metadata.id,
            "type_name": metadata.type_name,
            "created_at": metadata.created_at,
            "updated_at": metadata.updated_at,
            "tags": metadata.tags,
            "data": obj.model_dump_json(),
        }

    def _deserialize(self, row: tuple) -> StoredObject[BaseModel] | None:
        """Deserialize a database row to a StoredObject."""
        id_, type_name, created_at, updated_at, tags, data = row

        obj_type = self._object_types.get(type_name)
        if obj_type is None:
            # Type not registered, cannot deserialize
            return None

        try:
            if isinstance(data, str):
                obj = obj_type.model_validate_json(data)
            else:
                obj = obj_type.model_validate(data)

            metadata = ObjectMetadata(
                id=id_,
                created_at=created_at,
                updated_at=updated_at,
                tags=tags if tags else [],
                type_name=type_name,
            )
            return StoredObject(metadata=metadata, obj=obj)
        except Exception:
            return None

    def put(
        self,
        obj: T,
        id: str | None = None,
        tags: list[str] | None = None,
    ) -> StoredObject[T]:
        """Store or update an object.

        If an object with the same ID exists, it will be updated.
        The updated_at timestamp is automatically set to now.

        Args:
            obj: The BaseModel object to store.
            id: Optional explicit ID. If not provided, uses id_field value or auto-generates UUID.
            tags: Optional initial tags for the object.

        Returns:
            StoredObject wrapping the object with its metadata.
        """
        # Register the type if not already registered
        obj_type = type(obj)
        if obj_type.__name__ not in self._object_types:
            self._object_types[obj_type.__name__] = obj_type

        # Check if this ID already exists (for update vs insert)
        obj_id = self._get_id_for_object(obj, id)
        existing = self._get_metadata(obj_id)

        if existing:
            # Update existing - preserve created_at, update updated_at
            metadata = ObjectMetadata(
                id=obj_id,
                created_at=existing.created_at,
                updated_at=existing.updated_at,
                tags=existing.tags if tags is None else tags,
                type_name=obj_type.__name__,
            )
            metadata.touch()
        else:
            # New object
            metadata = ObjectMetadata(
                id=obj_id,
                tags=tags if tags else [],
                type_name=obj_type.__name__,
            )

        data = self._serialize(obj, metadata)

        self._conn.execute(
            f"""
            INSERT OR REPLACE INTO {self.name} (id, type_name, created_at, updated_at, tags, data)
            VALUES (?, ?, ?, ?, ?, ?)
        """,
            [data["id"], data["type_name"], data["created_at"], data["updated_at"], data["tags"], data["data"]],
        )

        return StoredObject(metadata=metadata, obj=obj)

    def _get_metadata(self, obj_id: str) -> ObjectMetadata | None:
        """Get just the metadata for an object (internal use)."""
        result = self._conn.execute(
            f"SELECT id, type_name, created_at, updated_at, tags FROM {self.name} WHERE id = ?", [obj_id]
        ).fetchone()

        if result is None:
            return None

        id_, type_name, created_at, updated_at, tags = result
        return ObjectMetadata(
            id=id_,
            created_at=created_at,
            updated_at=updated_at,
            tags=tags if tags else [],
            type_name=type_name,
        )

    def get(self, obj_id: str) -> BaseModel | None:
        """Retrieve object by ID.

        Returns the object, or None if not found.
        Use get_stored() if you need the full StoredObject with metadata.
        """
        stored = self.get_stored(obj_id)
        return stored.obj if stored else None

    def get_stored(self, obj_id: str) -> StoredObject[BaseModel] | None:
        """Retrieve object by ID with metadata.

        Returns StoredObject containing the object and its metadata,
        or None if object not found.
        """
        result = self._conn.execute(
            f"SELECT id, type_name, created_at, updated_at, tags, data FROM {self.name} WHERE id = ?", [obj_id]
        ).fetchone()

        if result is None:
            return None

        return self._deserialize(result)

    def delete(self, obj_id: str) -> bool:
        """Delete object by ID.

        Returns True if object was deleted, False if not found.
        """
        result = self._conn.execute(f"DELETE FROM {self.name} WHERE id = ? RETURNING id", [obj_id]).fetchone()

        return result is not None

    def search(self, **criteria: Any) -> list[BaseModel]:
        """Search objects by field values.

        Searches within the JSON data field for matching values.
        Only supports exact matches on top-level fields.

        Example:
            pit.search(name="Alice", age=30)

        Returns:
            List of objects matching the criteria.
            Use search_stored() if you need the full StoredObject with metadata.
        """
        return [stored.obj for stored in self.search_stored(**criteria)]

    def search_stored(self, **criteria: Any) -> list[StoredObject[BaseModel]]:
        """Search objects by field values with metadata.

        Searches within the JSON data field for matching values.
        Only supports exact matches on top-level fields.

        Example:
            pit.search_stored(name="Alice", age=30)

        Returns:
            List of StoredObject instances matching the criteria.
        """
        if not criteria:
            return self.all_stored()

        conditions = []
        params = []

        for field, value in criteria.items():
            if field in ("id", "type_name", "created_at", "updated_at"):
                conditions.append(f"{field} = ?")
                params.append(value)
            elif field == "tags":
                # Special handling for tags - search for objects with specific tags
                if isinstance(value, list):
                    for tag in value:
                        conditions.append("list_contains(tags, ?)")
                        params.append(tag)
                else:
                    conditions.append("list_contains(tags, ?)")
                    params.append(value)
            else:
                # Search in JSON data
                if isinstance(value, str):
                    conditions.append(f"json_extract_string(data, '$.{field}') = ?")
                elif isinstance(value, bool):
                    conditions.append(f"json_extract(data, '$.{field}') = ?")
                    value = "true" if value else "false"
                elif isinstance(value, (int, float)):
                    conditions.append(f"CAST(json_extract(data, '$.{field}') AS DOUBLE) = ?")
                else:
                    conditions.append(f"json_extract(data, '$.{field}') = ?")
                params.append(value)

        where_clause = " AND ".join(conditions)
        query = f"SELECT id, type_name, created_at, updated_at, tags, data FROM {self.name} WHERE {where_clause}"

        results = self._conn.execute(query, params).fetchall()
        return [stored for row in results if (stored := self._deserialize(row)) is not None]

    def search_tags(self, tags: list[str], match_all: bool = False) -> list[BaseModel]:
        """Search objects by tags.

        Args:
            tags: List of tags to search for.
            match_all: If True, object must have all tags. If False, any tag matches.

        Returns:
            List of objects matching the tag criteria.
            Use search_tags_stored() if you need the full StoredObject with metadata.
        """
        return [stored.obj for stored in self.search_tags_stored(tags, match_all)]

    def search_tags_stored(self, tags: list[str], match_all: bool = False) -> list[StoredObject[BaseModel]]:
        """Search objects by tags with metadata.

        Args:
            tags: List of tags to search for.
            match_all: If True, object must have all tags. If False, any tag matches.

        Returns:
            List of StoredObject instances matching the tag criteria.
        """
        if not tags:
            return []

        if match_all:
            conditions = ["list_contains(tags, ?)" for _ in tags]
            where_clause = " AND ".join(conditions)
        else:
            conditions = ["list_contains(tags, ?)" for _ in tags]
            where_clause = " OR ".join(conditions)

        query = f"SELECT id, type_name, created_at, updated_at, tags, data FROM {self.name} WHERE {where_clause}"

        results = self._conn.execute(query, tags).fetchall()
        return [stored for row in results if (stored := self._deserialize(row)) is not None]

    def all(self) -> list[BaseModel]:
        """Return all objects in pit.

        Use all_stored() if you need the full StoredObject with metadata.
        """
        return [stored.obj for stored in self.all_stored()]

    def all_stored(self) -> list[StoredObject[BaseModel]]:
        """Return all objects in pit with metadata."""
        results = self._conn.execute(
            f"SELECT id, type_name, created_at, updated_at, tags, data FROM {self.name}"
        ).fetchall()
        return [stored for row in results if (stored := self._deserialize(row)) is not None]

    # Tag management methods
    def add_tag(self, obj_id: str, tag: str) -> bool:
        """Add a tag to an object's metadata.

        Args:
            obj_id: ID of the object.
            tag: Tag to add.

        Returns:
            True if tag was added, False if object not found or tag already exists.
        """
        metadata = self._get_metadata(obj_id)
        if metadata is None:
            return False

        if tag in metadata.tags:
            return False

        new_tags = metadata.tags + [tag]
        self._conn.execute(
            f"UPDATE {self.name} SET tags = ?, updated_at = ? WHERE id = ?",
            [new_tags, datetime.now(), obj_id],
        )
        return True

    def remove_tag(self, obj_id: str, tag: str) -> bool:
        """Remove a tag from an object's metadata.

        Args:
            obj_id: ID of the object.
            tag: Tag to remove.

        Returns:
            True if tag was removed, False if object not found or tag doesn't exist.
        """
        metadata = self._get_metadata(obj_id)
        if metadata is None:
            return False

        if tag not in metadata.tags:
            return False

        new_tags = [t for t in metadata.tags if t != tag]
        self._conn.execute(
            f"UPDATE {self.name} SET tags = ?, updated_at = ? WHERE id = ?",
            [new_tags, datetime.now(), obj_id],
        )
        return True

    def get_tags(self, obj_id: str) -> list[str] | None:
        """Get tags for an object.

        Args:
            obj_id: ID of the object.

        Returns:
            List of tags, or None if object not found.
        """
        metadata = self._get_metadata(obj_id)
        if metadata is None:
            return None
        return metadata.tags

    def count(self) -> int:
        """Return count of objects in pit."""
        result = self._conn.execute(f"SELECT COUNT(*) FROM {self.name}").fetchone()
        return result[0] if result else 0

    def clear(self) -> int:
        """Remove all objects from pit. Returns count of deleted objects."""
        count = self.count()
        self._conn.execute(f"DELETE FROM {self.name}")
        return count

    def close(self) -> None:
        """Close the database connection."""
        self._conn.close()

    def __enter__(self) -> "Pit":
        """Context manager entry."""
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Context manager exit - closes database connection."""
        self.close()

    def __len__(self) -> int:
        """Return count of objects."""
        return self.count()

    def __contains__(self, obj_id: str) -> bool:
        """Check if object with ID exists in pit."""
        result = self._conn.execute(f"SELECT 1 FROM {self.name} WHERE id = ?", [obj_id]).fetchone()
        return result is not None
