"""Workspace usage balance - included allowance + purchased credits."""
from __future__ import annotations
from datetime import datetime, timezone
from typing import Optional
from uuid import UUID

from src.domain.admin.repository import SubscriptionPlanRepository
from src.domain.billing.repository import (
    SubscriptionRepository,
    WorkspaceCreditLedgerRepository,
    WorkspaceUsageAllowanceRepository,
)


class WorkspaceUsageBalanceResult:
    """Resolved balance for a workspace."""

    def __init__(
        self,
        included_allowance_remaining: int,
        purchased_topup_remaining: int,
        total_available: int,
        has_active_subscription: bool,
        plan_slug: Optional[str] = None,
    ):
        self.included_allowance_remaining = included_allowance_remaining
        self.purchased_topup_remaining = purchased_topup_remaining
        self.total_available = total_available
        self.has_active_subscription = has_active_subscription
        self.plan_slug = plan_slug


class WorkspaceUsageBalanceService:
    """
    Resolves workspace usage balance:
    - Included plan allowance (consumed first)
    - Purchased top-up credits (consumed second)
    """

    def __init__(
        self,
        subscription_repo: SubscriptionRepository,
        plan_repo: SubscriptionPlanRepository,
        allowance_repo: WorkspaceUsageAllowanceRepository,
        ledger_repo: WorkspaceCreditLedgerRepository,
    ):
        self._subscription_repo = subscription_repo
        self._plan_repo = plan_repo
        self._allowance_repo = allowance_repo
        self._ledger_repo = ledger_repo

    async def get_balance(self, workspace_id: UUID) -> WorkspaceUsageBalanceResult:
        """Return total available units and breakdown."""
        sub = await self._subscription_repo.get_by_workspace(workspace_id)
        included_remaining = 0
        plan_slug = None

        if sub:
            plan_slug = sub.plan_slug
            allowance = await self._allowance_repo.get_current_for_workspace(workspace_id)
            if allowance:
                included_remaining = allowance.remaining_units
            elif sub.plan_slug:
                # No allowance record yet - plan may include units; create allowance if period active
                plan = await self._plan_repo.get_by_slug(sub.plan_slug)
                if plan and plan.limits and sub.current_period_start and sub.current_period_end:
                    included = plan.limits.get("included_ai_units", 0) or plan.limits.get("ai_minutes", 0)
                    if included:
                        ps = sub.current_period_start
                        pe = sub.current_period_end
                        period_start = (
                            datetime.fromtimestamp(ps, tz=timezone.utc)
                            if isinstance(ps, (int, float))
                            else ps
                        )
                        period_end = (
                            datetime.fromtimestamp(pe, tz=timezone.utc)
                            if isinstance(pe, (int, float))
                            else pe
                        )
                        allowance = await self._allowance_repo.create(
                            workspace_id=workspace_id,
                            period_start=period_start,
                            period_end=period_end,
                            included_units=included,
                            subscription_id=sub.id,
                        )
                        included_remaining = allowance.remaining_units
        else:
            # Free plan: allowance may exist without subscription (subscription_id=null)
            allowance = await self._allowance_repo.get_current_for_workspace(workspace_id)
            if allowance:
                included_remaining = allowance.remaining_units
                plan_slug = "free"

        purchased_remaining = await self._ledger_repo.sum_purchased_balance(workspace_id)
        total = included_remaining + purchased_remaining

        return WorkspaceUsageBalanceResult(
            included_allowance_remaining=included_remaining,
            purchased_topup_remaining=purchased_remaining,
            total_available=total,
            has_active_subscription=sub is not None,
            plan_slug=plan_slug,
        )

    async def can_consume(self, workspace_id: UUID, units: int) -> bool:
        """True if workspace has at least `units` available."""
        balance = await self.get_balance(workspace_id)
        return balance.total_available >= units

    async def consume_units(
        self,
        workspace_id: UUID,
        units: int,
        *,
        source_type: str,
        source_id: Optional[str] = None,
        project_id: Optional[UUID] = None,
        codegen_job_id: Optional[UUID] = None,
        description: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> bool:
        """
        Consume units: allowance first, then purchased top-up.
        Uses advisory lock + row locking for concurrency safety.

        Transaction contract: MUST be called inside an outer DB transaction.
        Does not open its own transaction. Lock released when transaction commits.
        Returns True if consumption succeeded.
        """
        if units <= 0:
            return True

        # Serialize consumption per workspace; lock released when transaction commits
        await self._ledger_repo.acquire_consumption_lock(workspace_id)

        # Re-read balance after lock (consistent snapshot)
        allowance = await self._allowance_repo.get_current_for_workspace_for_update(workspace_id)
        included_remaining = allowance.remaining_units if allowance else 0
        purchased_remaining = await self._ledger_repo.sum_purchased_balance(workspace_id)
        total_available = included_remaining + purchased_remaining
        if total_available < units:
            return False

        remaining_to_consume = units

        # Consume from included allowance first (allowance row already locked)
        if included_remaining > 0 and allowance:
            from_allowance = min(remaining_to_consume, allowance.remaining_units)
            if from_allowance > 0:
                await self._allowance_repo.consume_units(allowance.id, from_allowance)
                await self._ledger_repo.add_entry(
                    workspace_id=workspace_id,
                    entry_type="consumption",
                    units_delta=-from_allowance,
                    source_bucket="subscription_allowance",
                    source_type=source_type,
                    source_id=source_id,
                    subscription_id=allowance.subscription_id,
                    project_id=project_id,
                    codegen_job_id=codegen_job_id,
                    description=description,
                    metadata={**(metadata or {}), "allowance_id": str(allowance.id)},
                )
                remaining_to_consume -= from_allowance

        # Consume from purchased top-up
        if remaining_to_consume > 0:
            balance_after = await self._ledger_repo.sum_purchased_balance(workspace_id) - remaining_to_consume
            await self._ledger_repo.add_entry(
                workspace_id=workspace_id,
                entry_type="consumption",
                units_delta=-remaining_to_consume,
                balance_after=balance_after,
                source_bucket="purchased_topup",
                source_type=source_type,
                source_id=source_id,
                project_id=project_id,
                codegen_job_id=codegen_job_id,
                description=description,
                metadata=metadata,
            )

        return True