"""
KenobiX - High-Performance Document Database

A SQLite3-backed document database with proper indexing for 15-665x faster operations.

Based on KenobiDB by Harrison Erd (https://github.com/patx/kenobi)
Enhanced with SQLite3 JSON optimizations and generated column indexes.

Key features:
1. Generated columns with indexes for specified fields (15-53x faster searches)
2. Automatic index usage in queries
3. Better concurrency model (no RLock for reads)
4. Cursor-based pagination option
5. Query plan analysis tools
6. 80-665x faster update operations

Copyright (c) 2025 KenobiX Contributors
Original KenobiDB Copyright (c) Harrison Erd
Licensed under BSD-3-Clause
"""

from __future__ import annotations

import contextlib
import json
import re
import sqlite3
from concurrent.futures import ThreadPoolExecutor
from threading import RLock
from typing import Any


class KenobiX:
    """
    KenobiX - High-performance document database with SQLite3 JSON optimization.

    Performance improvements over basic document stores:
    - 15-53x faster searches on indexed fields
    - 80-665x faster update operations
    - Minimal storage overhead (VIRTUAL generated columns)
    - Automatic index usage with fallback to json_extract

    Example:
        # Specify fields to index at creation
        db = KenobiX('test.db', indexed_fields=['name', 'age', 'city'])

        # Or add indexes dynamically
        db.create_index('email')

        # Queries automatically use indexes when available
        db.search('name', 'John')  # Uses index - 50x faster!
        db.search('tags', 'python')  # No index - falls back to json_extract
    """

    def __init__(self, file: str, indexed_fields: list[str] | None = None):
        """
        Initialize the database with optional field indexing.

        Args:
            file: Path to SQLite database file
            indexed_fields: List of document fields to create indexes for
                          Example: ['name', 'age', 'email']
        """
        self.file = file
        self._write_lock = RLock()  # Only for writes
        self._connection = sqlite3.connect(self.file, check_same_thread=False)
        self._indexed_fields: set[str] = set(indexed_fields or [])
        self.executor = ThreadPoolExecutor(max_workers=5)

        # Transaction state
        self._in_transaction = False
        self._savepoint_counter = 0

        self._add_regexp_support(self._connection)
        self._initialize_db()

    def _initialize_db(self):
        """Create table with generated columns for indexed fields."""
        with self._write_lock:
            # Build CREATE TABLE with generated columns
            columns = ["id INTEGER PRIMARY KEY AUTOINCREMENT", "data TEXT NOT NULL"]

            # Add generated columns for indexed fields (skip 'id' since it's the primary key)
            for field in self._indexed_fields:
                # Skip 'id' field - it's already the primary key
                if field == "id":
                    continue

                # Use VIRTUAL to avoid storage overhead
                # STORED would duplicate data but might be faster for some queries
                safe_field = self._sanitize_field_name(field)
                columns.append(
                    f"{safe_field} TEXT GENERATED ALWAYS AS "
                    f"(json_extract(data, '$.{field}')) VIRTUAL"
                )

            create_table = (
                f"CREATE TABLE IF NOT EXISTS documents (\n    {', '.join(columns)}\n)"
            )
            self._connection.execute(create_table)

            # Create indexes on generated columns
            for field in self._indexed_fields:
                safe_field = self._sanitize_field_name(field)
                self._connection.execute(
                    f"CREATE INDEX IF NOT EXISTS idx_{safe_field} ON documents({safe_field})"
                )

            self._connection.execute("PRAGMA journal_mode=WAL")
            self._connection.commit()

    @staticmethod
    def _sanitize_field_name(field: str) -> str:
        """Convert field name to valid SQL identifier."""
        # Remove special characters, replace with underscore
        return "".join(c if c.isalnum() else "_" for c in field)

    @staticmethod
    def _add_regexp_support(conn):
        """Add REGEXP function support."""

        def regexp(pattern, value):
            return re.search(pattern, value) is not None

        conn.create_function("REGEXP", 2, regexp)

    def create_index(self, field: str) -> bool:
        """
        Dynamically create an index on a field.

        Note: This requires recreating the table. Better to specify
        indexed_fields at initialization for production use.

        Args:
            field: Document field to index (e.g., 'email')

        Returns:
            True if index was created
        """
        if field in self._indexed_fields:
            return False  # Already indexed

        with self._write_lock:
            # For existing database, need to:
            # 1. Create new table with additional generated column
            # 2. Copy data
            # 3. Drop old table
            # 4. Rename new table
            # (Simplified here - production would need careful migration)

            self._indexed_fields.add(field)
            safe_field = self._sanitize_field_name(field)

            # Try to add column (SQLite limitation: can't add generated columns to existing tables easily)
            # In real implementation, would need table recreation or schema versioning
            try:
                self._connection.execute(
                    f"ALTER TABLE documents ADD COLUMN {safe_field} TEXT "
                    f"GENERATED ALWAYS AS (json_extract(data, '$.{field}')) VIRTUAL"
                )
                self._connection.execute(
                    f"CREATE INDEX idx_{safe_field} ON documents({safe_field})"
                )
                self._maybe_commit()
                return True
            except sqlite3.OperationalError:
                # Column already exists or can't be added
                return False

    def insert(self, document: dict[str, Any]) -> int:
        """
        Insert a document. Uses write lock.

        Args:
            document: Dictionary to insert

        Returns:
            The ID of the inserted document

        Raises:
            TypeError: If document is not a dict
        """
        if not isinstance(document, dict):
            msg = "Must insert a dict"
            raise TypeError(msg)

        with self._write_lock:
            cursor = self._connection.execute(
                "INSERT INTO documents (data) VALUES (?)", (json.dumps(document),)
            )
            self._maybe_commit()
            # SQLite always returns lastrowid after INSERT
            assert cursor.lastrowid is not None
            return cursor.lastrowid

    def insert_many(self, document_list: list[dict[str, Any]]) -> list[int]:
        """
        Insert multiple documents (list of dicts) into the database.

        Args:
            document_list: The list of documents to insert.

        Returns:
            List of IDs of the inserted documents.

        Raises:
            TypeError: If the provided object is not a list of dicts.
        """
        if not isinstance(document_list, list) or not all(
            isinstance(doc, dict) for doc in document_list
        ):
            msg = "Must insert a list of dicts"
            raise TypeError(msg)

        with self._write_lock:
            cursor = self._connection.execute("SELECT MAX(id) FROM documents")
            last_id = cursor.fetchone()[0] or 0

            self._connection.executemany(
                "INSERT INTO documents (data) VALUES (?)",
                [(json.dumps(doc),) for doc in document_list],
            )
            self._maybe_commit()

            # Return the range of IDs that were inserted
            return list(range(last_id + 1, last_id + 1 + len(document_list)))

    def remove(self, key: str, value: Any) -> int:
        """
        Remove all documents where the given key matches the specified value.

        Args:
            key: The field name to match.
            value: The value to match.

        Returns:
            Number of documents removed.

        Raises:
            ValueError: If 'key' is empty or 'value' is None.
        """
        if not key or not isinstance(key, str):
            msg = "key must be a non-empty string"
            raise ValueError(msg)
        if value is None:
            msg = "value cannot be None"
            raise ValueError(msg)

        with self._write_lock:
            if key in self._indexed_fields:
                safe_field = self._sanitize_field_name(key)
                query = f"DELETE FROM documents WHERE {safe_field} = ?"
                result = self._connection.execute(query, (value,))
            else:
                query = "DELETE FROM documents WHERE json_extract(data, '$.' || ?) = ?"
                result = self._connection.execute(query, (key, value))
            self._maybe_commit()
            return result.rowcount

    def update(self, id_key: str, id_value: Any, new_dict: dict[str, Any]) -> bool:
        """
        Update documents that match (id_key == id_value) by merging new_dict.

        Args:
            id_key: The field name to match.
            id_value: The value to match.
            new_dict: A dictionary of changes to apply.

        Returns:
            True if at least one document was updated, False otherwise.

        Raises:
            TypeError: If new_dict is not a dict.
            ValueError: If id_key is invalid or id_value is None.
        """
        if not isinstance(new_dict, dict):
            msg = "new_dict must be a dictionary"
            raise TypeError(msg)
        if not id_key or not isinstance(id_key, str):
            msg = "id_key must be a non-empty string"
            raise ValueError(msg)
        if id_value is None:
            msg = "id_value cannot be None"
            raise ValueError(msg)

        with self._write_lock:
            if id_key in self._indexed_fields:
                safe_field = self._sanitize_field_name(id_key)
                select_query = f"SELECT data FROM documents WHERE {safe_field} = ?"
                update_query = f"UPDATE documents SET data = ? WHERE {safe_field} = ?"
                cursor = self._connection.execute(select_query, (id_value,))
            else:
                select_query = (
                    "SELECT data FROM documents WHERE json_extract(data, '$.' || ?) = ?"
                )
                update_query = "UPDATE documents SET data = ? WHERE json_extract(data, '$.' || ?) = ?"
                cursor = self._connection.execute(select_query, (id_key, id_value))

            documents = cursor.fetchall()
            if not documents:
                return False

            for row in documents:
                document = json.loads(row[0])
                if not isinstance(document, dict):
                    continue
                document.update(new_dict)

                if id_key in self._indexed_fields:
                    self._connection.execute(
                        update_query, (json.dumps(document), id_value)
                    )
                else:
                    self._connection.execute(
                        update_query, (json.dumps(document), id_key, id_value)
                    )

            self._maybe_commit()
            return True

    def purge(self) -> bool:
        """
        Remove all documents from the database.

        Returns:
            True upon successful purge.
        """
        with self._write_lock:
            self._connection.execute("DELETE FROM documents")
            self._maybe_commit()
            return True

    def search(
        self, key: str, value: Any, limit: int = 100, offset: int = 0
    ) -> list[dict]:
        """
        Search documents. No lock needed for reads!
        Automatically uses index if available.

        Args:
            key: Field name to search
            value: Value to match
            limit: Max results to return
            offset: Skip this many results

        Returns:
            List of matching documents
        """
        if not key or not isinstance(key, str):
            msg = "Key must be a non-empty string"
            raise ValueError(msg)

        # Special case: searching by 'id' uses the primary key column directly
        # since we can't create a generated column for it
        if key == "id":
            # Search using json_extract on the id field in the data
            query = (
                "SELECT data FROM documents "
                "WHERE json_extract(data, '$.id') = ? "
                "LIMIT ? OFFSET ?"
            )
            cursor = self._connection.execute(query, (value, limit, offset))
        # Check if field is indexed - if so, use direct column query
        elif key in self._indexed_fields:
            safe_field = self._sanitize_field_name(key)
            query = (
                f"SELECT data FROM documents WHERE {safe_field} = ? LIMIT ? OFFSET ?"
            )
            cursor = self._connection.execute(query, (value, limit, offset))
        else:
            # Fall back to json_extract (no index)
            query = (
                "SELECT data FROM documents "
                "WHERE json_extract(data, '$.' || ?) = ? "
                "LIMIT ? OFFSET ?"
            )
            cursor = self._connection.execute(query, (key, value, limit, offset))

        return [json.loads(row[0]) for row in cursor.fetchall()]

    def search_optimized(self, **filters) -> list[dict]:
        """
        Multi-field search with automatic index usage.

        Example:
            db.search_optimized(name='John', age=30, city='NYC')
            # If all fields indexed: 70x faster than separate searches

        Args:
            **filters: field=value pairs to search

        Returns:
            List of matching documents
        """
        if not filters:
            return self.all()

        # Build WHERE clause using indexed columns when possible
        where_parts: list[str] = []
        params: list[Any] = []

        for key, value in filters.items():
            if key in self._indexed_fields:
                safe_field = self._sanitize_field_name(key)
                where_parts.append(f"{safe_field} = ?")
            else:
                where_parts.append(f"json_extract(data, '$.{key}') = ?")
            params.append(value)

        where_clause = " AND ".join(where_parts)
        query = f"SELECT data FROM documents WHERE {where_clause}"

        cursor = self._connection.execute(query, params)
        return [json.loads(row[0]) for row in cursor.fetchall()]

    def all(self, limit: int = 100, offset: int = 0) -> list[dict]:
        """Get all documents. No read lock needed."""
        query = "SELECT data FROM documents LIMIT ? OFFSET ?"
        cursor = self._connection.execute(query, (limit, offset))
        return [json.loads(row[0]) for row in cursor.fetchall()]

    def all_cursor(self, after_id: int | None = None, limit: int = 100) -> dict:
        """
        Cursor-based pagination for better performance on large datasets.

        Args:
            after_id: Continue from this document ID
            limit: Max results to return

        Returns:
            Dict with 'documents', 'next_cursor', 'has_more'

        Example:
            # First page
            result = db.all_cursor(limit=100)
            documents = result['documents']

            # Next page
            if result['has_more']:
                result = db.all_cursor(after_id=result['next_cursor'], limit=100)
        """
        if after_id:
            query = "SELECT id, data FROM documents WHERE id > ? ORDER BY id LIMIT ?"
            cursor = self._connection.execute(query, (after_id, limit + 1))
        else:
            query = "SELECT id, data FROM documents ORDER BY id LIMIT ?"
            cursor = self._connection.execute(query, (limit + 1,))

        rows = cursor.fetchall()
        has_more = len(rows) > limit

        if has_more:
            rows = rows[:limit]

        documents = [json.loads(row[1]) for row in rows]
        next_cursor = rows[-1][0] if rows else None

        return {
            "documents": documents,
            "next_cursor": next_cursor,
            "has_more": has_more,
        }

    def search_pattern(
        self, key: str, pattern: str, limit: int = 100, offset: int = 0
    ) -> list[dict]:
        """
        Search documents matching a regex pattern.

        Args:
            key: The document field to match on.
            pattern: The regex pattern to match.
            limit: The maximum number of documents to return.
            offset: The starting point for retrieval.

        Returns:
            List of matching documents (dicts).

        Raises:
            ValueError: If the key or pattern is invalid.
        """
        if not key or not isinstance(key, str):
            msg = "key must be a non-empty string"
            raise ValueError(msg)
        if not pattern or not isinstance(pattern, str):
            msg = "pattern must be a non-empty string"
            raise ValueError(msg)

        # Pattern search can't use indexes regardless
        query = """
            SELECT data FROM documents
            WHERE json_extract(data, '$.' || ?) REGEXP ?
            LIMIT ? OFFSET ?
        """
        cursor = self._connection.execute(query, (key, pattern, limit, offset))
        return [json.loads(row[0]) for row in cursor.fetchall()]

    def find_any(self, key: str, value_list: list[Any]) -> list[dict]:
        """
        Return documents where key matches any value in value_list.

        Args:
            key: The document field to match on.
            value_list: A list of possible values.

        Returns:
            A list of matching documents.
        """
        if not value_list:
            return []

        placeholders = ", ".join(["?"] * len(value_list))

        # Use indexed column if available, otherwise json_each
        if key in self._indexed_fields:
            safe_field = self._sanitize_field_name(key)
            query = f"""
                SELECT DISTINCT data
                FROM documents
                WHERE {safe_field} IN ({placeholders})
            """
            cursor = self._connection.execute(query, value_list)
        else:
            query = f"""
                SELECT DISTINCT documents.data
                FROM documents, json_each(documents.data, '$.' || ?)
                WHERE json_each.value IN ({placeholders})
            """
            cursor = self._connection.execute(query, [key] + value_list)

        return [json.loads(row[0]) for row in cursor.fetchall()]

    def find_all(self, key: str, value_list: list[Any]) -> list[dict]:
        """
        Return documents where the key contains all values in value_list.

        Args:
            key: The field to match.
            value_list: The required values to match.

        Returns:
            A list of matching documents.
        """
        if not value_list:
            return []

        placeholders = ", ".join(["?"] * len(value_list))

        # This requires json_each to check array membership
        query = f"""
            SELECT documents.data
            FROM documents
            WHERE (
                SELECT COUNT(DISTINCT value)
                FROM json_each(documents.data, '$.' || ?)
                WHERE value IN ({placeholders})
            ) = ?
        """
        cursor = self._connection.execute(query, [key] + value_list + [len(value_list)])
        return [json.loads(row[0]) for row in cursor.fetchall()]

    def execute_async(self, func, *args, **kwargs):
        """
        Execute a function asynchronously using a thread pool.

        Args:
            func: The function to execute.
            *args: Arguments for the function.
            **kwargs: Keyword arguments for the function.

        Returns:
            concurrent.futures.Future: A Future object representing the execution.
        """
        return self.executor.submit(func, *args, **kwargs)

    def explain(self, operation: str, *args) -> list[tuple]:
        """
        Show query execution plan for optimization.

        Example:
            plan = db.explain('search', 'name', 'John')
            for row in plan:
                print(row)
            # Output: SEARCH documents USING INDEX idx_name (name=?)
            #    or:  SCAN documents  (bad - no index!)

        Args:
            operation: Method name ('search', 'all', etc.)
            *args: Arguments to the method

        Returns:
            List of query plan tuples from EXPLAIN QUERY PLAN
        """
        # Map operation to SQL query
        if operation == "search":
            key, value = args[0], args[1]
            if key in self._indexed_fields:
                safe_field = self._sanitize_field_name(key)
                query = f"EXPLAIN QUERY PLAN SELECT data FROM documents WHERE {safe_field} = ?"
                cursor = self._connection.execute(query, (value,))
            else:
                query = "EXPLAIN QUERY PLAN SELECT data FROM documents WHERE json_extract(data, '$.' || ?) = ?"
                cursor = self._connection.execute(query, (key, value))
        elif operation == "all":
            query = "EXPLAIN QUERY PLAN SELECT data FROM documents"
            cursor = self._connection.execute(query)
        else:
            msg = f"Unknown operation: {operation}"
            raise ValueError(msg)

        return cursor.fetchall()

    def get_indexed_fields(self) -> set[str]:
        """Return set of fields that have indexes."""
        return self._indexed_fields.copy()

    def stats(self) -> dict[str, Any]:
        """
        Get database statistics.

        Returns:
            Dict with document count, database size, indexed fields, etc.
        """
        cursor = self._connection.execute("SELECT COUNT(*) FROM documents")
        doc_count = cursor.fetchone()[0]

        cursor = self._connection.execute(
            "SELECT page_count * page_size FROM pragma_page_count(), pragma_page_size()"
        )
        db_size = cursor.fetchone()[0]

        return {
            "document_count": doc_count,
            "database_size_bytes": db_size,
            "indexed_fields": list(self._indexed_fields),
            "wal_mode": True,
        }

    def _maybe_commit(self):
        """Commit if not in a transaction."""
        if not self._in_transaction:
            self._connection.commit()

    def begin(self):
        """
        Begin a new transaction.

        Allows multiple operations to be grouped atomically.
        All operations after begin() will not be committed until commit() is called.

        Example:
            db.begin()
            try:
                db.insert({"user": "Alice"})
                db.insert({"user": "Bob"})
                db.commit()  # Both succeed or both fail
            except:
                db.rollback()

        Raises:
            RuntimeError: If already in a transaction
        """
        if self._in_transaction:
            msg = "Already in a transaction. Use savepoint() for nested transactions."
            raise RuntimeError(msg)

        with self._write_lock:
            self._connection.execute("BEGIN")
            self._in_transaction = True

    def commit(self):
        """
        Commit the current transaction.

        Makes all changes since begin() permanent.

        Raises:
            RuntimeError: If not in a transaction
        """
        if not self._in_transaction:
            msg = "Not in a transaction"
            raise RuntimeError(msg)

        with self._write_lock:
            self._connection.commit()
            self._in_transaction = False
            self._savepoint_counter = 0

    def rollback(self):
        """
        Rollback the current transaction.

        Discards all changes since begin().

        Raises:
            RuntimeError: If not in a transaction
        """
        if not self._in_transaction:
            msg = "Not in a transaction"
            raise RuntimeError(msg)

        with self._write_lock:
            self._connection.rollback()
            self._in_transaction = False
            self._savepoint_counter = 0

    def savepoint(self, name: str | None = None) -> str:
        """
        Create a savepoint within a transaction.

        Savepoints allow partial rollback within a transaction.

        Args:
            name: Optional savepoint name (auto-generated if not provided)

        Returns:
            Savepoint name

        Example:
            db.begin()
            db.insert({"user": "Alice"})
            sp = db.savepoint()
            db.insert({"user": "Bob"})
            db.rollback_to(sp)  # Rolls back Bob, keeps Alice
            db.commit()

        Raises:
            RuntimeError: If not in a transaction
        """
        if not self._in_transaction:
            msg = "Must be in a transaction to create savepoint"
            raise RuntimeError(msg)

        if name is None:
            self._savepoint_counter += 1
            name = f"sp_{self._savepoint_counter}"

        with self._write_lock:
            self._connection.execute(f"SAVEPOINT {name}")

        return name

    def rollback_to(self, savepoint: str):
        """
        Rollback to a specific savepoint.

        Args:
            savepoint: Savepoint name (from savepoint() method)

        Raises:
            RuntimeError: If not in a transaction
        """
        if not self._in_transaction:
            msg = "Not in a transaction"
            raise RuntimeError(msg)

        with self._write_lock:
            self._connection.execute(f"ROLLBACK TO SAVEPOINT {savepoint}")

    def release_savepoint(self, savepoint: str):
        """
        Release a savepoint (commit it within the transaction).

        Args:
            savepoint: Savepoint name

        Raises:
            RuntimeError: If not in a transaction
        """
        if not self._in_transaction:
            msg = "Not in a transaction"
            raise RuntimeError(msg)

        with self._write_lock:
            self._connection.execute(f"RELEASE SAVEPOINT {savepoint}")

    def transaction(self):
        """
        Context manager for transactions.

        Automatically begins transaction on enter and commits on exit.
        Rolls back on exception.

        Example:
            with db.transaction():
                db.insert({"user": "Alice"})
                db.insert({"user": "Bob"})
                # Both committed together

        Returns:
            Transaction context manager
        """
        return Transaction(self)

    def close(self):
        """Shutdown executor and close connection."""
        self.executor.shutdown()
        with self._write_lock:
            self._connection.close()


class Transaction:
    """
    Context manager for database transactions.

    Provides automatic transaction management with commit/rollback.
    """

    def __init__(self, db: KenobiX):
        """
        Initialize transaction context manager.

        Args:
            db: KenobiX database instance
        """
        self.db = db
        self._savepoint: str | None = None

    def __enter__(self):
        """Begin transaction or create savepoint if nested."""
        if self.db._in_transaction:
            # Nested transaction - use savepoint
            self._savepoint = self.db.savepoint()
        else:
            # Top-level transaction
            self.db.begin()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Commit on success, rollback on exception."""
        try:
            if exc_type is not None:
                # Exception occurred - rollback
                if self._savepoint:
                    self.db.rollback_to(self._savepoint)
                else:
                    self.db.rollback()
                return False  # Re-raise exception
            # Success - commit
            if self._savepoint:
                self.db.release_savepoint(self._savepoint)
            else:
                self.db.commit()
            return True
        except sqlite3.Error:
            # Error during commit/rollback - ensure we rollback
            if not self._savepoint and self.db._in_transaction:
                with contextlib.suppress(sqlite3.Error):
                    self.db.rollback()
            raise
