from __future__ import annotations

import logging
import uuid
from typing import TYPE_CHECKING

import attrs

from codeflash_api.auth.models import (
    APIKey,
    Organization,
    Subscription,
)

if TYPE_CHECKING:
    import asyncpg  # type: ignore[import-untyped]

log = logging.getLogger(__name__)


@attrs.define
class Queries:
    """
    Raw SQL queries against the asyncpg connection pool.
    """

    _pool: asyncpg.Pool

    async def get_api_key_by_hash(self, hashed_key: str) -> APIKey | None:
        """
        Look up an API key by its SHA-384 hash.
        """
        row = await self._pool.fetchrow(
            "SELECT id, key, suffix, name, user_id, tier,"
            " organization_id, created_at, last_used"
            " FROM cf_api_keys WHERE key = $1",
            hashed_key,
        )
        if row is None:
            return None
        return APIKey(**dict(row))

    async def update_api_key_last_used(self, api_key_id: int) -> None:
        """
        Update the ``last_used`` timestamp for an API key.
        """
        await self._pool.execute(
            "UPDATE cf_api_keys SET last_used = NOW() WHERE id = $1",
            api_key_id,
        )

    async def get_organization(
        self, organization_id: str
    ) -> Organization | None:
        """
        Fetch an organization by ID.
        """
        row = await self._pool.fetchrow(
            "SELECT id, name, github_org_id, description,"
            " website, auto_add_github_members, added_by,"
            " privacy_mode, subscription, created_at,"
            " updated_at"
            " FROM organizations WHERE id = $1",
            organization_id,
        )
        if row is None:
            return None
        return Organization(**dict(row))

    async def get_subscription(self, user_id: str) -> Subscription | None:
        """
        Fetch a subscription by user ID.
        """
        row = await self._pool.fetchrow(
            "SELECT id, user_id, plan_type,"
            " optimizations_used,"
            " total_lifetime_optimizations,"
            " optimizations_limit,"
            " subscription_status,"
            " stripe_customer_id,"
            " stripe_subscription_id,"
            " current_period_start,"
            " current_period_end,"
            " cancel_at_period_end,"
            " cancellation_request_date,"
            " created_at, updated_at"
            " FROM subscriptions WHERE user_id = $1",
            user_id,
        )
        if row is None:
            return None
        return Subscription(**dict(row))

    async def create_subscription(
        self,
        user_id: str,
        plan_type: str,
        optimizations_limit: int | str,
    ) -> Subscription:
        """
        Create a new subscription with defaults.
        """
        sub_id = str(uuid.uuid4())
        await self._pool.execute(
            "INSERT INTO subscriptions"
            " (id, user_id, plan_type,"
            "  optimizations_limit,"
            "  subscription_status,"
            "  optimizations_used)"
            " VALUES ($1, $2, $3, $4, $5, $6)",
            sub_id,
            user_id,
            plan_type,
            int(optimizations_limit),
            "active",
            0,
        )
        return Subscription(
            id=sub_id,
            user_id=user_id,
            plan_type=plan_type,
            optimizations_limit=int(optimizations_limit),
        )

    async def increment_usage(self, user_id: str, cost: int) -> None:
        """
        Atomically increment usage counters for a subscription.
        """
        await self._pool.execute(
            "UPDATE subscriptions"
            " SET optimizations_used ="
            "     optimizations_used + $1,"
            "     total_lifetime_optimizations ="
            "     total_lifetime_optimizations + $1"
            " WHERE user_id = $2",
            cost,
            user_id,
        )
