import asyncio
import dataclasses
from typing import Any, Dict, List, Tuple, Optional

# =============================================================================
# EXCEPTIONS
# =============================================================================

class DatabaseError(Exception):
    """Base class for database errors."""
    pass

class PoolExhaustedError(DatabaseError):
    """Raised when the connection pool cannot provide a connection within timeout."""
    pass

class TransactionError(DatabaseError):
    """Raised when a transaction fails and is rolled back."""
    pass

class ValidationError(ValueError):
    """Raised when input parameters fail validation."""
    pass

# =============================================================================
# MOCK DATABASE LAYER
# =============================================================================

@dataclasses.dataclass
class User:
    id: int
    name: str
    label: str

@dataclasses.dataclass
class Post:
    id: int
    user_id: int
    content: str

class MockDatabaseState:
    """Holds the actual 'in-memory' data."""
    def __init__(self):
        self.users: Dict[int, User] = {
            i: User(id=i, name=f"User_{i}", label="default") for i in range(1, 26)
        }
        self.posts: List[Post] = []
        for u_id in self.users:
            for p_idx in range(5):
                self.posts.append(Post(id=len(self.posts) + 1, user_id=u_id, content=f"Post {p_idx} by {u_id}"))

class MockConnection:
    """Simulates a single database connection."""
    def __init__(self, pool: 'MockPool', state: MockDatabaseState):
        self.pool = pool
        self.state = state
        self._transaction_buffer: List[Tuple[int, str]] = []  # Stores (user_id, new_label)
        self._in_transaction = False

    async def execute(self, query: str, params: Tuple[Any, ...]) -> int:
        """Simulates executing a command. Returns number of rows affected."""
        # Simulate network latency
        await asyncio.sleep(0.01)

        # Parameterized Query Check: 
        # In a real driver, the query string contains placeholders (?), not values.
        # We simulate an error if someone tries to pass a query that looks like it was interpolated.
        if "'" in query and not any(isinstance(p, str) for p in params):
            # This is a naive check to demonstrate the concept of preventing injection
            pass 

        if "UPDATE users SET label =" in query:
            new_label, user_id = params
            if user_id not in self.state.users:
                return 0  # No rows affected (user doesn't exist)
            
            if self._in_transaction:
                self._transaction_buffer.append((user_id, new_label))
            else:
                # Autocommit mode
                self.state.users[user_id].label = new_label
            return 1

        return 0

    async def fetch(self, query: str, params: Tuple[Any, ...]) -> List[Dict[str, Any]]:
        """Simulates fetching rows."""
        await asyncio.sleep(0.01)

        if "SELECT * FROM users" in query:
            # Handle Pagination logic via params (LIMIT, OFFSET)
            # Note: In a real app, we'd parse the SQL. Here we simulate based on params.
            if "WHERE id =" in query:
                u_id = params[0]
                user = self.state.users.get(u_id)
                return [{"id": user.id, "name": user.name, "label": user.label}] if user else []
            
            # Pagination: params might be (limit, offset) or just empty
            limit = params[0] if len(params) > 0 else 10
            offset = params[1] if len(params) > 1 else 0
            
            all_users = list(self.state.users.values())
            sliced = all_users[offset : offset + limit]
            return [{"id": u.id, "name": u.name, "label": u.label} for u in sliced]

        if "SELECT COUNT(*) FROM users" in query:
            return [{"count": len(self.state.users)}]

        if "SELECT * FROM posts WHERE user_id =" in query:
            u_id = params[0]
            user_posts = [p for p in self.state.posts if p.user_id == u_id]
            return [{"id": p.id, "content": p.content} for p in user_posts]

        return []

    async def begin(self):
        self._in_transaction = True
        self._transaction_buffer = []

    async def commit(self):
        if not self._in_transaction:
            return
        # Apply buffered updates to the real state
        for user_id, new_label in self._transaction_buffer:
            self.state.users[user_id].label = new_label
        self._in_transaction = False
        self._transaction_buffer = []

    async def rollback(self):
        self._transaction_buffer = []
        self._in_transaction = False

    async def release(self):
        await self.pool._release_connection(self)

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.release()

class MockPool:
    """Simulates a connection pool with a fixed size."""
    def __init__(self, max_size: int, state: MockDatabaseState):
        self.max_size = max_size
        self.state = state
        self._semaphore = asyncio.Semaphore(max_size)
        self._in_use_count = 0

    async def acquire(self) -> MockConnection:
        try:
            # Requirement 5: Respect max-wait (2s)
            await asyncio.wait_for(self._semaphore.acquire(), timeout=2.0)
            self._in_use_count += 1
            return MockConnection(self, self.state)
        except asyncio.TimeoutError:
            raise PoolExhaustedError("Could not acquire connection: Pool exhausted.")

    async def _release_connection(self, conn: MockConnection):
        self._in_use_count -= 1
        self._semaphore.release()

    @property
    def in_use(self) -> int:
        return self._in_use_count

# =============================================================================
# DATA ACCESS SERVICE
# =============================================================================

class UserService:
    def __init__(self, pool: MockPool):
        self.pool = pool

    async def get_users(self, page: int, page_size: int) -> Dict[str, Any]:
        """Requirement 1 & 3: Pagination and Connection Discipline."""
        if page < 1 or not (1 <= page_size <= 100):
            raise ValidationError("Invalid page or page_size. Page >= 1, 1 <= size <= 100.")

        async with await self.pool.acquire() as conn:
            # Get total count
            count_res = await conn.fetch("SELECT COUNT(*) FROM users", ())
            total = count_res[0]["count"]

            # Calculate offset
            offset = (page - 1) * page_size
            
            # Get items
            items = await conn.fetch("SELECT * FROM users LIMIT ? OFFSET ?", (page_size, offset))

            return {
                "items": items,
                "page": page,
                "page_size": page_size,
                "total": total
            }

    async def get_user_with_posts(self, user_id: int) -> Dict[str, Any]:
        """Requirement 2: Efficient fetching (no N+1)."""
        async with await self.pool.acquire() as conn:
            user_res = await conn.fetch("SELECT * FROM users WHERE id = ?", (user_id,))
            if not user_res:
                raise KeyError(f"User with id {user_id} not found.")
            
            user = user_res[0]
            # Fetch posts in the same connection context
            posts = await conn.fetch("SELECT * FROM posts WHERE user_id = ?", (user_id,))
            
            return {**user, "posts": posts}

    async def relabel_users(self, pairs: List[Tuple[int, str]]) -> int:
        """Requirement 4: Transactional bulk update."""
        async with await self.pool.acquire() as conn:
            await conn.begin()
            try:
                total_updated = 0
                for user_id, new_label in pairs:
                    # Requirement 2: Parameterized queries (passing params as tuple)
                    rows_affected = await conn.execute(
                        "UPDATE users SET label = ? WHERE id = ?", 
                        (new_label, user_id)
                    )
                    if rows_affected == 0:
                        # If any update fails (user doesn't exist), trigger rollback
                        raise TransactionError(f"Update failed: User {user_id} not found.")
                    total_updated += rows_affected
                
                await conn.commit()
                return total_updated
            except Exception as e:
                await conn.rollback()
                # Re-raise or return 0 based on requirement. We raise to be explicit.
                raise e

# =============================================================================
# TEST SUITE
# =============================================================================

async def run_tests():
    print("🚀 Starting Tests...\n")
    state = MockDatabaseState()
    pool = MockPool(max_size=5, state=state)
    service = UserService(pool)

    # --- Test A: Connection Leak Check ---
    print("Test A: Checking for connection leaks (100 sequential calls)...")
    for _ in range(100):
        await service.get_users(page=1, page_size=10)
    assert pool.in_use == 0, f"Leak detected! Connections in use: {pool.in_use}"
    print("✅ No leaks detected.")

    # --- Test B: Pagination Math ---
    print("\nTest B: Verifying pagination math...")
    # Total users is 25. Page size 10.
    # Page 1: 1-10, Page 2: 11-20, Page 3: 21-25
    res_p1 = await service.get_users(page=1, page_size=10)
    assert res_p1["total"] == 25
    assert len(res_p1["items"]) == 10

    res_p3 = await service.get_users(page=3, page_size=10)
    assert len(res_p3["items"]) == 5  # Remainder
    
    res_out = await service.get_users(page=10, page_size=10)
    assert len(res_out["items"]) == 0 # Out of range
    print("✅ Pagination math is correct.")

    # --- Test C: Transactional Integrity (Rollback) ---
    print("\nTest C: Verifying transactional rollback...")
    # Initial state check for user 1
    user_1_before = (await service.get_user_with_posts(1))["label"]
    
    # Attempt bulk update: User 2 is valid, User 999 is invalid.
    # This should cause the whole batch to fail.
    try:
        await service.relabel_users([(2, "new_label"), (999, "fail_label")])
    except TransactionError:
        pass # Expected
    
    user_1_after = (await service.get_user_with_posts(1))["label"]
    user_2_after = (await service.get_user_with_posts(2))["label"]
    
    assert user_1_after == user_1_before, "User 1 changed despite transaction failure!"
    assert user_2_after == "default", "User 2 changed despite transaction failure!"
    print("✅ Transaction rolled back successfully. No partial updates applied.")

    # --- Test D: Concurrency & Pool Exhaustion ---
    print("\nTest D: Verifying concurrency (10 concurrent requests on pool of 5)...")
    # We trigger 10 tasks. Since max_size is 5, some will wait for the semaphore.
    # This tests that they don't crash and eventually complete.
    tasks = [service.get_user_with_posts(i) for i in range(1, 11)]
    results = await asyncio.gather(*tasks)
    assert len(results) == 10
    assert pool.in_use == 0
    print("✅ Concurrent requests completed successfully.")

    # --- Test E: Parameterized Query / Injection Simulation ---
    print("\nTest E: Verifying parameterized query usage...")
    # We check if the service correctly handles a "malicious" string by passing it as a param
    # rather than interpolating it. 
    malicious_label = "'; DROP TABLE users; --"
    await service.relabel_users([(5, malicious_label)])
    user_5 = await service.get_user_with_posts(5)
    assert user_5["label"] == malicious_label, "The label was not treated as a literal string!"
    print("✅ Parameterized query logic verified.")

    print("\n✨ ALL TESTS PASSED! ✨")

if __name__ == "__main__":
    try:
        asyncio.run(run_tests())
    except KeyboardInterrupt:
        pass