from __future__ import annotations
import logging
from uuid import UUID
from typing import Optional, Dict, Any
from typing import List

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.exc import IntegrityError
from sqlalchemy import text

from ..models import Account
from ..models import UsageEvent

log = logging.getLogger(__name__)


class SimpleUsageTracker:
    """A simple service to track credit usage for billing."""

    def __init__(self, db_session: AsyncSession):
        self.db = db_session

    async def track_usage(
        self,
        account_id: UUID,
        credits: int,
        event_type: str = "api_call",
        metadata: Optional[Dict[str, Any]] = None,
    ):
        """
        Tracks a usage event by saving it to the database.
        Includes an idempotency key to prevent double-billing.
        """
        if credits <= 0:
            return

        # Generate idempotency key to prevent duplicate events
        metadata = metadata or {}
        job_id = metadata.get("job_id")
        idempotency_key = f"{account_id}:{event_type}:{job_id}" if job_id else None

        try:
            event = UsageEvent(
                account_id=account_id,
                credits_consumed=credits,
                event_type=event_type,
                idempotency_key=idempotency_key,
                metadata=metadata,
            )
            self.db.add(event)
            await self.db.commit()
            log.debug(
                f"Tracked usage for account {account_id}: {credits} credits for {event_type}"
            )
        except IntegrityError:
            # This happens if the idempotency_key already exists. It's not an error.
            await self.db.rollback()
            log.warning(
                f"Duplicate usage event detected and ignored for key: {idempotency_key}"
            )
        except Exception as e:
            await self.db.rollback()
            log.error(
                f"Failed to track usage for account {account_id}: {e}", exc_info=True
            )

    async def track_usage_batch(self, events: List[Dict[str, Any]]):
        """Tracks multiple usage events efficiently in a single transaction."""
        if not events:
            return

        usage_events = [UsageEvent(**event) for event in events]
        try:
            self.db.add_all(usage_events)
            await self.db.commit()
            log.debug(f"Batch-tracked {len(usage_events)} usage events.")
        except Exception as e:
            await self.db.rollback()
            log.error(f"Failed to batch-track usage events: {e}", exc_info=True)

    async def get_current_usage(self, account_id: UUID) -> int:
        """
        Gets the current calendar month's total credit usage for an account
        from the pre-aggregated summary view.
        """
        try:
            # Query the continuous aggregate for the current month
            query = text("""
                SELECT total_credits
                FROM monthly_billing_summary
                WHERE account_id = :account_id AND month = date_trunc('month', NOW());
            """)
            result = await self.db.execute(query, {"account_id": account_id})
            usage = result.scalar_one_or_none()
            return usage or 0
        except Exception as e:
            log.error(
                f"Could not get current usage for account {account_id}: {e}",
                exc_info=True,
            )
            return 0

    async def check_credit_limit(self, account_id: UUID, credits_needed: int) -> bool:
        """
        Checks if an account has enough credits for an operation.
        Returns True if the account has sufficient credits, False otherwise.
        """
        account = await self.db.get(Account, account_id)
        if not account:
            log.warning(f"Account not found for credit check: {account_id}")
            return False  # Default to fail-safe

        return account.credits_balance >= credits_needed
