import re
import json
from uuid import uuid4
from datetime import datetime, timezone
from typing import (
    Any,
    Dict,
    Iterable,
    List,
    Optional,
    AsyncGenerator,
    Union,
    Tuple,
    Set,
)
from pathlib import Path
import aiosqlite
import unittest
from types import SimpleNamespace
import asyncio
import aiohttp
from aiohttp import web
import socket
import os
import atexit


class AsyncDataSet:
    """
    Distributed AsyncDataSet with client-server model over Unix sockets.

    Parameters:
    - file: Path to the SQLite database file
    - socket_path: Path to Unix socket for inter-process communication (default: "ads.sock")
    - max_concurrent_queries: Maximum concurrent database operations (default: 1)
      Set to 1 to prevent "database is locked" errors with SQLite
    """

    _KV_TABLE = "__kv_store"
    _DEFAULT_COLUMNS = {
        "uid": "TEXT PRIMARY KEY",
        "created_at": "TEXT",
        "updated_at": "TEXT",
        "deleted_at": "TEXT",
    }

    def __init__(
        self, file: str, socket_path: str = "ads.sock", max_concurrent_queries: int = 1
    ):
        self._file = file
        self._socket_path = socket_path
        self._table_columns_cache: Dict[str, Set[str]] = {}
        self._is_server = None  # None means not initialized yet
        self._server = None
        self._runner = None
        self._client_session = None
        self._initialized = False
        self._init_lock = None  # Will be created when needed
        self._max_concurrent_queries = max_concurrent_queries
        self._db_semaphore = None  # Will be created when needed
        self._db_connection = None  # Persistent database connection

    async def _ensure_initialized(self):
        """Ensure the instance is initialized as server or client."""
        if self._initialized:
            return

        # Create lock if needed (first time in async context)
        if self._init_lock is None:
            self._init_lock = asyncio.Lock()

        async with self._init_lock:
            if self._initialized:
                return

            await self._initialize()

            # Create semaphore for database operations (server only)
            if self._is_server:
                self._db_semaphore = asyncio.Semaphore(self._max_concurrent_queries)

            self._initialized = True

    async def _initialize(self):
        """Initialize as server or client."""
        try:
            # Try to create Unix socket
            if os.path.exists(self._socket_path):
                # Check if socket is active
                sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
                try:
                    sock.connect(self._socket_path)
                    sock.close()
                    # Socket is active, we're a client
                    self._is_server = False
                    await self._setup_client()
                except (ConnectionRefusedError, FileNotFoundError):
                    # Socket file exists but not active, remove and become server
                    os.unlink(self._socket_path)
                    self._is_server = True
                    await self._setup_server()
            else:
                # No socket file, become server
                self._is_server = True
                await self._setup_server()
        except Exception as e:
            # Fallback to client mode
            self._is_server = False
            await self._setup_client()

        # Establish persistent database connection
        self._db_connection = await aiosqlite.connect(self._file)

    async def _setup_server(self):
        """Setup aiohttp server."""
        app = web.Application()

        # Add routes for all methods
        app.router.add_post("/insert", self._handle_insert)
        app.router.add_post("/update", self._handle_update)
        app.router.add_post("/delete", self._handle_delete)
        app.router.add_post("/upsert", self._handle_upsert)
        app.router.add_post("/get", self._handle_get)
        app.router.add_post("/find", self._handle_find)
        app.router.add_post("/count", self._handle_count)
        app.router.add_post("/exists", self._handle_exists)
        app.router.add_post("/kv_set", self._handle_kv_set)
        app.router.add_post("/kv_get", self._handle_kv_get)
        app.router.add_post("/execute_raw", self._handle_execute_raw)
        app.router.add_post("/query_raw", self._handle_query_raw)
        app.router.add_post("/query_one", self._handle_query_one)
        app.router.add_post("/create_table", self._handle_create_table)
        app.router.add_post("/insert_unique", self._handle_insert_unique)
        app.router.add_post("/aggregate", self._handle_aggregate)

        self._runner = web.AppRunner(app)
        await self._runner.setup()
        self._server = web.UnixSite(self._runner, self._socket_path)
        await self._server.start()

        # Register cleanup
        def cleanup():
            if os.path.exists(self._socket_path):
                os.unlink(self._socket_path)

        atexit.register(cleanup)

    async def _setup_client(self):
        """Setup aiohttp client."""
        connector = aiohttp.UnixConnector(path=self._socket_path)
        self._client_session = aiohttp.ClientSession(connector=connector)

    async def _make_request(self, endpoint: str, data: Dict[str, Any]) -> Any:
        """Make HTTP request to server."""
        if not self._client_session:
            await self._setup_client()

        url = f"http://localhost/{endpoint}"
        async with self._client_session.post(url, json=data) as resp:
            result = await resp.json()
            if result.get("error"):
                raise Exception(result["error"])
            return result.get("result")

    # Server handlers
    async def _handle_insert(self, request):
        data = await request.json()
        try:
            result = await self._server_insert(
                data["table"], data["args"], data.get("return_id", False)
            )
            return web.json_response({"result": result})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_update(self, request):
        data = await request.json()
        try:
            result = await self._server_update(
                data["table"], data["args"], data.get("where")
            )
            return web.json_response({"result": result})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_delete(self, request):
        data = await request.json()
        try:
            result = await self._server_delete(data["table"], data.get("where"))
            return web.json_response({"result": result})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_upsert(self, request):
        data = await request.json()
        try:
            result = await self._server_upsert(
                data["table"], data["args"], data.get("where")
            )
            return web.json_response({"result": result})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_get(self, request):
        data = await request.json()
        try:
            result = await self._server_get(data["table"], data.get("where"))
            return web.json_response({"result": result})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_find(self, request):
        data = await request.json()
        try:
            result = await self._server_find(
                data["table"],
                data.get("where"),
                limit=data.get("limit", 0),
                offset=data.get("offset", 0),
                order_by=data.get("order_by"),
            )
            return web.json_response({"result": result})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_count(self, request):
        data = await request.json()
        try:
            result = await self._server_count(data["table"], data.get("where"))
            return web.json_response({"result": result})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_exists(self, request):
        data = await request.json()
        try:
            result = await self._server_exists(data["table"], data["where"])
            return web.json_response({"result": result})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_kv_set(self, request):
        data = await request.json()
        try:
            await self._server_kv_set(
                data["key"], data["value"], table=data.get("table")
            )
            return web.json_response({"result": None})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_kv_get(self, request):
        data = await request.json()
        try:
            result = await self._server_kv_get(
                data["key"], default=data.get("default"), table=data.get("table")
            )
            return web.json_response({"result": result})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_execute_raw(self, request):
        data = await request.json()
        try:
            result = await self._server_execute_raw(
                data["sql"],
                tuple(data.get("params", [])) if data.get("params") else None,
            )
            return web.json_response(
                {"result": result.rowcount if hasattr(result, "rowcount") else None}
            )
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_query_raw(self, request):
        data = await request.json()
        try:
            result = await self._server_query_raw(
                data["sql"],
                tuple(data.get("params", [])) if data.get("params") else None,
            )
            return web.json_response({"result": result})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_query_one(self, request):
        data = await request.json()
        try:
            result = await self._server_query_one(
                data["sql"],
                tuple(data.get("params", [])) if data.get("params") else None,
            )
            return web.json_response({"result": result})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_create_table(self, request):
        data = await request.json()
        try:
            await self._server_create_table(
                data["table"], data["schema"], data.get("constraints")
            )
            return web.json_response({"result": None})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_insert_unique(self, request):
        data = await request.json()
        try:
            result = await self._server_insert_unique(
                data["table"], data["args"], data["unique_fields"]
            )
            return web.json_response({"result": result})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    async def _handle_aggregate(self, request):
        data = await request.json()
        try:
            result = await self._server_aggregate(
                data["table"],
                data["function"],
                data.get("column", "*"),
                data.get("where"),
            )
            return web.json_response({"result": result})
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)

    # Public methods that delegate to server or client
    async def insert(
        self, table: str, args: Dict[str, Any], return_id: bool = False
    ) -> Union[str, int]:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_insert(table, args, return_id)
        else:
            return await self._make_request(
                "insert", {"table": table, "args": args, "return_id": return_id}
            )

    async def update(
        self, table: str, args: Dict[str, Any], where: Optional[Dict[str, Any]] = None
    ) -> int:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_update(table, args, where)
        else:
            return await self._make_request(
                "update", {"table": table, "args": args, "where": where}
            )

    async def delete(self, table: str, where: Optional[Dict[str, Any]] = None) -> int:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_delete(table, where)
        else:
            return await self._make_request("delete", {"table": table, "where": where})

    async def upsert(
        self, table: str, args: Dict[str, Any], where: Optional[Dict[str, Any]] = None
    ) -> str | None:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_upsert(table, args, where)
        else:
            return await self._make_request(
                "upsert", {"table": table, "args": args, "where": where}
            )

    async def get(
        self, table: str, where: Optional[Dict[str, Any]] = None
    ) -> Optional[Dict[str, Any]]:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_get(table, where)
        else:
            return await self._make_request("get", {"table": table, "where": where})

    async def find(
        self,
        table: str,
        where: Optional[Dict[str, Any]] = None,
        *,
        limit: int = 0,
        offset: int = 0,
        order_by: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_find(
                table, where, limit=limit, offset=offset, order_by=order_by
            )
        else:
            return await self._make_request(
                "find",
                {
                    "table": table,
                    "where": where,
                    "limit": limit,
                    "offset": offset,
                    "order_by": order_by,
                },
            )

    async def count(self, table: str, where: Optional[Dict[str, Any]] = None) -> int:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_count(table, where)
        else:
            return await self._make_request("count", {"table": table, "where": where})

    async def exists(self, table: str, where: Dict[str, Any]) -> bool:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_exists(table, where)
        else:
            return await self._make_request("exists", {"table": table, "where": where})

    async def kv_set(self, key: str, value: Any, *, table: str | None = None) -> None:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_kv_set(key, value, table=table)
        else:
            return await self._make_request(
                "kv_set", {"key": key, "value": value, "table": table}
            )

    async def kv_get(
        self, key: str, *, default: Any = None, table: str | None = None
    ) -> Any:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_kv_get(key, default=default, table=table)
        else:
            return await self._make_request(
                "kv_get", {"key": key, "default": default, "table": table}
            )

    async def execute_raw(self, sql: str, params: Optional[Tuple] = None) -> Any:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_execute_raw(sql, params)
        else:
            return await self._make_request(
                "execute_raw", {"sql": sql, "params": list(params) if params else None}
            )

    async def query_raw(
        self, sql: str, params: Optional[Tuple] = None
    ) -> List[Dict[str, Any]]:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_query_raw(sql, params)
        else:
            return await self._make_request(
                "query_raw", {"sql": sql, "params": list(params) if params else None}
            )

    async def query_one(
        self, sql: str, params: Optional[Tuple] = None
    ) -> Optional[Dict[str, Any]]:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_query_one(sql, params)
        else:
            return await self._make_request(
                "query_one", {"sql": sql, "params": list(params) if params else None}
            )

    async def create_table(
        self,
        table: str,
        schema: Dict[str, str],
        constraints: Optional[List[str]] = None,
    ):
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_create_table(table, schema, constraints)
        else:
            return await self._make_request(
                "create_table",
                {"table": table, "schema": schema, "constraints": constraints},
            )

    async def insert_unique(
        self, table: str, args: Dict[str, Any], unique_fields: List[str]
    ) -> Union[str, None]:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_insert_unique(table, args, unique_fields)
        else:
            return await self._make_request(
                "insert_unique",
                {"table": table, "args": args, "unique_fields": unique_fields},
            )

    async def aggregate(
        self,
        table: str,
        function: str,
        column: str = "*",
        where: Optional[Dict[str, Any]] = None,
    ) -> Any:
        await self._ensure_initialized()
        if self._is_server:
            return await self._server_aggregate(table, function, column, where)
        else:
            return await self._make_request(
                "aggregate",
                {
                    "table": table,
                    "function": function,
                    "column": column,
                    "where": where,
                },
            )

    async def transaction(self):
        """Context manager for transactions."""
        await self._ensure_initialized()
        if self._is_server:
            return TransactionContext(self._db_connection, self._db_semaphore)
        else:
            raise NotImplementedError("Transactions not supported in client mode")

    # Server implementation methods (original logic)
    @staticmethod
    def _utc_iso() -> str:
        return (
            datetime.now(timezone.utc)
            .replace(microsecond=0)
            .isoformat()
            .replace("+00:00", "Z")
        )

    @staticmethod
    def _py_to_sqlite_type(value: Any) -> str:
        if value is None:
            return "TEXT"
        if isinstance(value, bool):
            return "INTEGER"
        if isinstance(value, int):
            return "INTEGER"
        if isinstance(value, float):
            return "REAL"
        if isinstance(value, (bytes, bytearray, memoryview)):
            return "BLOB"
        return "TEXT"

    async def _get_table_columns(self, table: str) -> Set[str]:
        """Get actual columns that exist in the table."""
        if table in self._table_columns_cache:
            return self._table_columns_cache[table]

        columns = set()
        try:
            async with self._db_semaphore:
                async with self._db_connection.execute(
                    f"PRAGMA table_info({table})"
                ) as cursor:
                    async for row in cursor:
                        columns.add(row[1])  # Column name is at index 1
            self._table_columns_cache[table] = columns
        except:
            pass
        return columns

    async def _invalidate_column_cache(self, table: str):
        """Invalidate column cache for a table."""
        if table in self._table_columns_cache:
            del self._table_columns_cache[table]

    async def _ensure_column(self, table: str, name: str, value: Any) -> None:
        col_type = self._py_to_sqlite_type(value)
        try:
            async with self._db_semaphore:
                await self._db_connection.execute(
                    f"ALTER TABLE {table} ADD COLUMN `{name}` {col_type}"
                )
                await self._db_connection.commit()
            await self._invalidate_column_cache(table)
        except aiosqlite.OperationalError as e:
            if "duplicate column name" in str(e).lower():
                pass  # Column already exists
            else:
                raise

    async def _ensure_table(self, table: str, col_sources: Dict[str, Any]) -> None:
        # Always include default columns
        cols = self._DEFAULT_COLUMNS.copy()

        # Add columns from col_sources
        for key, val in col_sources.items():
            if key not in cols:
                cols[key] = self._py_to_sqlite_type(val)

        columns_sql = ", ".join(f"`{k}` {t}" for k, t in cols.items())
        async with self._db_semaphore:
            await self._db_connection.execute(
                f"CREATE TABLE IF NOT EXISTS {table} ({columns_sql})"
            )
            await self._db_connection.commit()
        await self._invalidate_column_cache(table)

    async def _table_exists(self, table: str) -> bool:
        """Check if a table exists."""
        async with self._db_semaphore:
            async with self._db_connection.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table,)
            ) as cursor:
                return await cursor.fetchone() is not None

    _RE_NO_COLUMN = re.compile(r"(?:no such column:|has no column named) (\w+)")
    _RE_NO_TABLE = re.compile(r"no such table: (\w+)")

    @classmethod
    def _missing_column_from_error(
        cls, err: aiosqlite.OperationalError
    ) -> Optional[str]:
        m = cls._RE_NO_COLUMN.search(str(err))
        return m.group(1) if m else None

    @classmethod
    def _missing_table_from_error(
        cls, err: aiosqlite.OperationalError
    ) -> Optional[str]:
        m = cls._RE_NO_TABLE.search(str(err))
        return m.group(1) if m else None

    async def _safe_execute(
        self,
        table: str,
        sql: str,
        params: Iterable[Any],
        col_sources: Dict[str, Any],
        max_retries: int = 10,
    ) -> aiosqlite.Cursor:
        retries = 0
        while retries < max_retries:
            try:
                async with self._db_semaphore:
                    cursor = await self._db_connection.execute(sql, params)
                    await self._db_connection.commit()
                    return cursor
            except aiosqlite.OperationalError as err:
                retries += 1
                err_str = str(err).lower()

                # Handle missing column
                col = self._missing_column_from_error(err)
                if col:
                    if col in col_sources:
                        await self._ensure_column(table, col, col_sources[col])
                    else:
                        # Column not in sources, ensure it with NULL/TEXT type
                        await self._ensure_column(table, col, None)
                    continue

                # Handle missing table
                tbl = self._missing_table_from_error(err)
                if tbl:
                    await self._ensure_table(tbl, col_sources)
                    continue

                # Handle other column-related errors
                if "has no column named" in err_str:
                    # Extract column name differently
                    match = re.search(r"table \w+ has no column named (\w+)", err_str)
                    if match:
                        col_name = match.group(1)
                        if col_name in col_sources:
                            await self._ensure_column(
                                table, col_name, col_sources[col_name]
                            )
                        else:
                            await self._ensure_column(table, col_name, None)
                        continue

                raise
        raise Exception(f"Max retries ({max_retries}) exceeded")

    async def _filter_existing_columns(
        self, table: str, data: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Filter data to only include columns that exist in the table."""
        if not await self._table_exists(table):
            return data

        existing_columns = await self._get_table_columns(table)
        if not existing_columns:
            return data

        return {k: v for k, v in data.items() if k in existing_columns}

    async def _safe_query(
        self,
        table: str,
        sql: str,
        params: Iterable[Any],
        col_sources: Dict[str, Any],
    ) -> AsyncGenerator[Dict[str, Any], None]:
        # Check if table exists first
        if not await self._table_exists(table):
            return

        max_retries = 10
        retries = 0

        while retries < max_retries:
            try:
                async with self._db_semaphore:
                    self._db_connection.row_factory = aiosqlite.Row
                    async with self._db_connection.execute(sql, params) as cursor:
                        # Fetch all rows while holding the semaphore
                        rows = await cursor.fetchall()

                # Yield rows after releasing the semaphore
                for row in rows:
                    yield dict(row)
                return
            except aiosqlite.OperationalError as err:
                retries += 1
                err_str = str(err).lower()

                # Handle missing table
                tbl = self._missing_table_from_error(err)
                if tbl:
                    # For queries, if table doesn't exist, just return empty
                    return

                # Handle missing column in WHERE clause or SELECT
                if "no such column" in err_str:
                    # For queries with missing columns, return empty
                    return

                raise

    @staticmethod
    def _build_where(where: Optional[Dict[str, Any]]) -> tuple[str, List[Any]]:
        if not where:
            return "", []
        clauses, vals = zip(*[(f"`{k}` = ?", v) for k, v in where.items()])
        return " WHERE " + " AND ".join(clauses), list(vals)

    async def _server_insert(
        self, table: str, args: Dict[str, Any], return_id: bool = False
    ) -> Union[str, int]:
        """Insert a record. If return_id=True, returns auto-incremented ID instead of UUID."""
        uid = str(uuid4())
        now = self._utc_iso()
        record = {
            "uid": uid,
            "created_at": now,
            "updated_at": now,
            "deleted_at": None,
            **args,
        }

        # Ensure table exists with all needed columns
        await self._ensure_table(table, record)

        # Handle auto-increment ID if requested
        if return_id and "id" not in args:
            # Ensure id column exists
            async with self._db_semaphore:
                # Add id column if it doesn't exist
                try:
                    await self._db_connection.execute(
                        f"ALTER TABLE {table} ADD COLUMN id INTEGER PRIMARY KEY AUTOINCREMENT"
                    )
                    await self._db_connection.commit()
                except aiosqlite.OperationalError as e:
                    if "duplicate column name" not in str(e).lower():
                        # Try without autoincrement constraint
                        try:
                            await self._db_connection.execute(
                                f"ALTER TABLE {table} ADD COLUMN id INTEGER"
                            )
                            await self._db_connection.commit()
                        except:
                            pass

            await self._invalidate_column_cache(table)

            # Insert and get lastrowid
            cols = "`" + "`, `".join(record.keys()) + "`"
            qs = ", ".join(["?"] * len(record))
            sql = f"INSERT INTO {table} ({cols}) VALUES ({qs})"
            cursor = await self._safe_execute(table, sql, list(record.values()), record)
            return cursor.lastrowid

        cols = "`" + "`, `".join(record) + "`"
        qs = ", ".join(["?"] * len(record))
        sql = f"INSERT INTO {table} ({cols}) VALUES ({qs})"
        await self._safe_execute(table, sql, list(record.values()), record)
        return uid

    async def _server_update(
        self,
        table: str,
        args: Dict[str, Any],
        where: Optional[Dict[str, Any]] = None,
    ) -> int:
        if not args:
            return 0

        # Check if table exists
        if not await self._table_exists(table):
            return 0

        args["updated_at"] = self._utc_iso()

        # Ensure all columns exist
        all_cols = {**args, **(where or {})}
        await self._ensure_table(table, all_cols)
        for col, val in all_cols.items():
            await self._ensure_column(table, col, val)

        set_clause = ", ".join(f"`{k}` = ?" for k in args)
        where_clause, where_params = self._build_where(where)
        sql = f"UPDATE {table} SET {set_clause}{where_clause}"
        params = list(args.values()) + where_params
        cur = await self._safe_execute(table, sql, params, all_cols)
        return cur.rowcount

    async def _server_delete(
        self, table: str, where: Optional[Dict[str, Any]] = None
    ) -> int:
        # Check if table exists
        if not await self._table_exists(table):
            return 0

        where_clause, where_params = self._build_where(where)
        sql = f"DELETE FROM {table}{where_clause}"
        cur = await self._safe_execute(table, sql, where_params, where or {})
        return cur.rowcount

    async def _server_upsert(
        self,
        table: str,
        args: Dict[str, Any],
        where: Optional[Dict[str, Any]] = None,
    ) -> str | None:
        if not args:
            raise ValueError("Nothing to update. Empty dict given.")
        args["updated_at"] = self._utc_iso()
        affected = await self._server_update(table, args, where)
        if affected:
            rec = await self._server_get(table, where)
            return rec.get("uid") if rec else None
        merged = {**(where or {}), **args}
        return await self._server_insert(table, merged)

    async def _server_get(
        self, table: str, where: Optional[Dict[str, Any]] = None
    ) -> Optional[Dict[str, Any]]:
        where_clause, where_params = self._build_where(where)
        sql = f"SELECT * FROM {table}{where_clause} LIMIT 1"
        async for row in self._safe_query(table, sql, where_params, where or {}):
            return row
        return None

    async def _server_find(
        self,
        table: str,
        where: Optional[Dict[str, Any]] = None,
        *,
        limit: int = 0,
        offset: int = 0,
        order_by: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """Find records with optional ordering."""
        where_clause, where_params = self._build_where(where)
        order_clause = f" ORDER BY {order_by}" if order_by else ""
        extra = (f" LIMIT {limit}" if limit else "") + (
            f" OFFSET {offset}" if offset else ""
        )
        sql = f"SELECT * FROM {table}{where_clause}{order_clause}{extra}"
        return [
            row async for row in self._safe_query(table, sql, where_params, where or {})
        ]

    async def _server_count(
        self, table: str, where: Optional[Dict[str, Any]] = None
    ) -> int:
        # Check if table exists
        if not await self._table_exists(table):
            return 0

        where_clause, where_params = self._build_where(where)
        sql = f"SELECT COUNT(*) FROM {table}{where_clause}"
        gen = self._safe_query(table, sql, where_params, where or {})
        async for row in gen:
            return next(iter(row.values()), 0)
        return 0

    async def _server_exists(self, table: str, where: Dict[str, Any]) -> bool:
        return (await self._server_count(table, where)) > 0

    async def _server_kv_set(
        self,
        key: str,
        value: Any,
        *,
        table: str | None = None,
    ) -> None:
        tbl = table or self._KV_TABLE
        json_val = json.dumps(value, default=str)
        await self._server_upsert(tbl, {"value": json_val}, {"key": key})

    async def _server_kv_get(
        self,
        key: str,
        *,
        default: Any = None,
        table: str | None = None,
    ) -> Any:
        tbl = table or self._KV_TABLE
        row = await self._server_get(tbl, {"key": key})
        if not row:
            return default
        try:
            return json.loads(row["value"])
        except Exception:
            return default

    async def _server_execute_raw(
        self, sql: str, params: Optional[Tuple] = None
    ) -> Any:
        """Execute raw SQL for complex queries like JOINs."""
        async with self._db_semaphore:
            cursor = await self._db_connection.execute(sql, params or ())
            await self._db_connection.commit()
            return cursor

    async def _server_query_raw(
        self, sql: str, params: Optional[Tuple] = None
    ) -> List[Dict[str, Any]]:
        """Execute raw SQL query and return results as list of dicts."""
        try:
            async with self._db_semaphore:
                self._db_connection.row_factory = aiosqlite.Row
                async with self._db_connection.execute(sql, params or ()) as cursor:
                    return [dict(row) async for row in cursor]
        except aiosqlite.OperationalError:
            # Return empty list if query fails
            return []

    async def _server_query_one(
        self, sql: str, params: Optional[Tuple] = None
    ) -> Optional[Dict[str, Any]]:
        """Execute raw SQL query and return single result."""
        results = await self._server_query_raw(sql + " LIMIT 1", params)
        return results[0] if results else None

    async def _server_create_table(
        self,
        table: str,
        schema: Dict[str, str],
        constraints: Optional[List[str]] = None,
    ):
        """Create table with custom schema and constraints. Always includes default columns."""
        # Merge default columns with custom schema
        full_schema = self._DEFAULT_COLUMNS.copy()
        full_schema.update(schema)

        columns = [f"`{col}` {dtype}" for col, dtype in full_schema.items()]
        if constraints:
            columns.extend(constraints)
        columns_sql = ", ".join(columns)

        async with self._db_semaphore:
            await self._db_connection.execute(
                f"CREATE TABLE IF NOT EXISTS {table} ({columns_sql})"
            )
            await self._db_connection.commit()
        await self._invalidate_column_cache(table)

    async def _server_insert_unique(
        self, table: str, args: Dict[str, Any], unique_fields: List[str]
    ) -> Union[str, None]:
        """Insert with unique constraint handling. Returns uid on success, None if duplicate."""
        try:
            return await self._server_insert(table, args)
        except aiosqlite.IntegrityError as e:
            if "UNIQUE" in str(e):
                return None
            raise

    async def _server_aggregate(
        self,
        table: str,
        function: str,
        column: str = "*",
        where: Optional[Dict[str, Any]] = None,
    ) -> Any:
        """Perform aggregate functions like SUM, AVG, MAX, MIN."""
        # Check if table exists
        if not await self._table_exists(table):
            return None

        where_clause, where_params = self._build_where(where)
        sql = f"SELECT {function}({column}) as result FROM {table}{where_clause}"
        result = await self._server_query_one(sql, tuple(where_params))
        return result["result"] if result else None

    async def close(self):
        """Close the connection or server."""
        if not self._initialized:
            return

        if self._client_session:
            await self._client_session.close()
        if self._runner:
            await self._runner.cleanup()
        if os.path.exists(self._socket_path) and self._is_server:
            os.unlink(self._socket_path)
        if self._db_connection:
            await self._db_connection.close()


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

    def __init__(
        self, db_connection: aiosqlite.Connection, semaphore: asyncio.Semaphore = None
    ):
        self.db_connection = db_connection
        self.semaphore = semaphore

    async def __aenter__(self):
        if self.semaphore:
            await self.semaphore.acquire()
        try:
            await self.db_connection.execute("BEGIN")
            return self.db_connection
        except:
            if self.semaphore:
                self.semaphore.release()
            raise

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        try:
            if exc_type is None:
                await self.db_connection.commit()
            else:
                await self.db_connection.rollback()
        finally:
            if self.semaphore:
                self.semaphore.release()


# Test cases remain the same but with additional tests for new functionality
class TestAsyncDataSet(unittest.IsolatedAsyncioTestCase):

    async def asyncSetUp(self):
        self.db_path = Path("temp_test.db")
        if self.db_path.exists():
            self.db_path.unlink()
        self.connector = AsyncDataSet(str(self.db_path), max_concurrent_queries=1)

    async def asyncTearDown(self):
        await self.connector.close()
        if self.db_path.exists():
            self.db_path.unlink()

    async def test_insert_and_get(self):
        await self.connector.insert("people", {"name": "John Doe", "age": 30})
        rec = await self.connector.get("people", {"name": "John Doe"})
        self.assertIsNotNone(rec)
        self.assertEqual(rec["name"], "John Doe")

    async def test_get_nonexistent(self):
        result = await self.connector.get("people", {"name": "Jane Doe"})
        self.assertIsNone(result)

    async def test_update(self):
        await self.connector.insert("people", {"name": "John Doe", "age": 30})
        await self.connector.update("people", {"age": 31}, {"name": "John Doe"})
        rec = await self.connector.get("people", {"name": "John Doe"})
        self.assertEqual(rec["age"], 31)

    async def test_order_by(self):
        await self.connector.insert("people", {"name": "Alice", "age": 25})
        await self.connector.insert("people", {"name": "Bob", "age": 30})
        await self.connector.insert("people", {"name": "Charlie", "age": 20})

        results = await self.connector.find("people", order_by="age ASC")
        self.assertEqual(results[0]["name"], "Charlie")
        self.assertEqual(results[-1]["name"], "Bob")

    async def test_raw_query(self):
        await self.connector.insert("people", {"name": "John", "age": 30})
        await self.connector.insert("people", {"name": "Jane", "age": 25})

        results = await self.connector.query_raw(
            "SELECT * FROM people WHERE age > ?", (26,)
        )
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0]["name"], "John")

    async def test_aggregate(self):
        await self.connector.insert("people", {"name": "John", "age": 30})
        await self.connector.insert("people", {"name": "Jane", "age": 25})
        await self.connector.insert("people", {"name": "Bob", "age": 35})

        avg_age = await self.connector.aggregate("people", "AVG", "age")
        self.assertEqual(avg_age, 30)

        max_age = await self.connector.aggregate("people", "MAX", "age")
        self.assertEqual(max_age, 35)

    async def test_insert_with_auto_id(self):
        # Test auto-increment ID functionality
        id1 = await self.connector.insert("posts", {"title": "First"}, return_id=True)
        id2 = await self.connector.insert("posts", {"title": "Second"}, return_id=True)
        self.assertEqual(id2, id1 + 1)

    async def test_transaction(self):
        await self.connector._ensure_initialized()
        if self.connector._is_server:
            async with self.connector.transaction() as conn:
                await conn.execute(
                    "INSERT INTO people (uid, name, age, created_at, updated_at) VALUES (?, ?, ?, ?, ?)",
                    ("test-uid", "John", 30, "2024-01-01", "2024-01-01"),
                )
                # Transaction will be committed

            rec = await self.connector.get("people", {"name": "John"})
            self.assertIsNotNone(rec)

    async def test_create_custom_table(self):
        schema = {
            "id": "INTEGER PRIMARY KEY AUTOINCREMENT",
            "username": "TEXT NOT NULL",
            "email": "TEXT NOT NULL",
            "score": "INTEGER DEFAULT 0",
        }
        constraints = ["UNIQUE(username)", "UNIQUE(email)"]

        await self.connector.create_table("users", schema, constraints)

        # Test that table was created with constraints
        result = await self.connector.insert_unique(
            "users",
            {"username": "john", "email": "john@example.com"},
            ["username", "email"],
        )
        self.assertIsNotNone(result)

        # Test duplicate insert
        result = await self.connector.insert_unique(
            "users",
            {"username": "john", "email": "different@example.com"},
            ["username", "email"],
        )
        self.assertIsNone(result)

    async def test_missing_table_operations(self):
        # Test operations on non-existent tables
        self.assertEqual(await self.connector.count("nonexistent"), 0)
        self.assertEqual(await self.connector.find("nonexistent"), [])
        self.assertIsNone(await self.connector.get("nonexistent"))
        self.assertFalse(await self.connector.exists("nonexistent", {"id": 1}))
        self.assertEqual(await self.connector.delete("nonexistent"), 0)
        self.assertEqual(
            await self.connector.update("nonexistent", {"name": "test"}), 0
        )

    async def test_auto_column_creation(self):
        # Insert with new columns that don't exist yet
        await self.connector.insert(
            "dynamic", {"col1": "value1", "col2": 42, "col3": 3.14}
        )

        # Add more columns in next insert
        await self.connector.insert(
            "dynamic", {"col1": "value2", "col4": True, "col5": None}
        )

        # All records should be retrievable
        records = await self.connector.find("dynamic")
        self.assertEqual(len(records), 2)


if __name__ == "__main__":
    unittest.main()