"""User billing store module."""

import asyncpg

from topix.datatypes.user_billing import UserBilling
from topix.store.postgres.pool import create_pool
from topix.store.postgres.user_billing import (
    create_user_billing,
    delete_user_billing_by_uid,
    get_user_billing_by_customer_id,
    get_user_billing_by_subscription_id,
    get_user_billing_by_uid,
    update_user_billing_by_uid,
    upsert_user_billing_by_uid,
)


class UserBillingStore:
    """Store for managing user billing in the database."""

    def __init__(self):
        """Initialize the UserBillingStore."""
        self._pg_pool: asyncpg.Pool | None = None
        self._owns_pool = False

    async def open(self, pool: asyncpg.Pool | None = None):
        """Open the store. Pass a shared pool, or omit to create a private one."""
        if pool is None:
            self._pg_pool = await create_pool()
            self._owns_pool = True
        else:
            self._pg_pool = pool
            self._owns_pool = False

    async def add_user_billing(self, billing: UserBilling):
        """Add a new user billing record."""
        async with self._pg_pool.acquire() as conn:
            await create_user_billing(conn, billing)

    async def get_user_billing(self, user_uid: str) -> UserBilling | None:
        """Retrieve user billing by user UID."""
        async with self._pg_pool.acquire() as conn:
            return await get_user_billing_by_uid(conn, user_uid)

    async def get_user_billing_by_stripe_customer_id(self, stripe_customer_id: str) -> UserBilling | None:
        """Retrieve user billing by Stripe customer ID."""
        async with self._pg_pool.acquire() as conn:
            return await get_user_billing_by_customer_id(conn, stripe_customer_id)

    async def get_user_billing_by_stripe_subscription_id(self, stripe_subscription_id: str) -> UserBilling | None:
        """Retrieve user billing by Stripe subscription ID."""
        async with self._pg_pool.acquire() as conn:
            return await get_user_billing_by_subscription_id(conn, stripe_subscription_id)

    async def update_user_billing(self, user_uid: str, data: dict):
        """Update user billing by user UID."""
        async with self._pg_pool.acquire() as conn:
            await update_user_billing_by_uid(conn, user_uid, data)

    async def upsert_user_billing(self, user_uid: str, data: dict):
        """Upsert user billing by user UID."""
        async with self._pg_pool.acquire() as conn:
            await upsert_user_billing_by_uid(conn, user_uid, data)

    async def delete_user_billing(self, user_uid: str):
        """Delete user billing by user UID."""
        async with self._pg_pool.acquire() as conn:
            await delete_user_billing_by_uid(conn, user_uid)

    async def close(self):
        """Close the store. Only closes the pool if this store created it."""
        if self._pg_pool and self._owns_pool:
            await self._pg_pool.close()
