from typing import Any, Dict, List, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from ..clients.base import DBClient


class QuerySet:
    """Rich result wrapper with chainable operations"""

    def __init__(self, rows: List[Dict], client: "DBClient", table: str, query=None):
        self.rows = rows or []
        self.client = client
        self.table = table
        self.query = query

    def all(self) -> List[Dict]:
        return self.rows

    def first(self) -> Optional[Dict]:
        return self.rows[0] if self.rows else None

    def one(self) -> Dict:
        if len(self.rows) != 1:
            raise ValueError(f"Expected exactly one row, got {len(self.rows)}")
        return self.rows[0]

    def count(self) -> int:
        return len(self.rows)

    async def update(self, **values: Any) -> int:
        if not self.rows:
            return 0
        ids = [r.get("id") for r in self.rows if "id" in r]
        if not ids:
            raise ValueError("Cannot perform bulk update: no 'id' field found")
        return await self.client._bulk_update(self.table, ids, values)

    async def remove(self) -> int:
        """
        Delete all rows represented by this QuerySet.

        Returns:
            Number of rows deleted

        Example:
            deleted = await result.remove()
        """
        return await self.delete()

    async def delete(self) -> int:
        if not self.rows:
            return 0
        ids = [r.get("id") for r in self.rows if "id" in r]
        if not ids:
            raise ValueError("Cannot perform bulk delete: no 'id' field found")
        return await self.client._bulk_delete(self.table, ids)
