from __future__ import annotations

from datetime import datetime
from typing import Optional
from uuid import UUID, uuid4

from sqlalchemy import func as sa_func, select, text
from sqlalchemy.ext.asyncio import AsyncSession

from src.domain.runtime_ai.models import RuntimeAiLedger, RuntimeAiPackage, RuntimeAiPurchase


class RuntimeAiPackageRepository:
    def __init__(self, db: AsyncSession) -> None:
        self._db = db

    async def get_by_slug(self, slug: str) -> Optional[RuntimeAiPackage]:
        result = await self._db.execute(select(RuntimeAiPackage).where(RuntimeAiPackage.slug == slug))
        return result.scalar_one_or_none()

    async def list_active(self) -> list[RuntimeAiPackage]:
        result = await self._db.execute(
            select(RuntimeAiPackage).where(RuntimeAiPackage.is_active.is_(True)).order_by(RuntimeAiPackage.sort_order)
        )
        return list(result.scalars().all())

    async def update(self, package_id: UUID, **fields) -> Optional[RuntimeAiPackage]:
        row = await self._db.get(RuntimeAiPackage, package_id)
        if not row:
            return None
        for key, value in fields.items():
            setattr(row, key, value)
        await self._db.flush()
        return row


class RuntimeAiPurchaseRepository:
    def __init__(self, db: AsyncSession) -> None:
        self._db = db

    async def get_by_session_id(self, session_id: str) -> Optional[RuntimeAiPurchase]:
        result = await self._db.execute(
            select(RuntimeAiPurchase).where(RuntimeAiPurchase.stripe_checkout_session_id == session_id)
        )
        return result.scalar_one_or_none()

    async def create(self, **fields) -> RuntimeAiPurchase:
        row = RuntimeAiPurchase(id=uuid4(), **fields)
        self._db.add(row)
        await self._db.flush()
        return row


class RuntimeAiLedgerRepository:
    def __init__(self, db: AsyncSession) -> None:
        self._db = db

    async def acquire_lock(self, workspace_id: UUID) -> None:
        await self._db.execute(
            text("SELECT pg_advisory_xact_lock(:ns, hashtext(:ws))"),
            {"ns": 42, "ws": f"runtime_ai:{workspace_id}"},
        )

    async def sum_balance(self, workspace_id: UUID) -> int:
        result = await self._db.execute(
            select(sa_func.coalesce(sa_func.sum(RuntimeAiLedger.millicredits_delta), 0)).where(
                RuntimeAiLedger.workspace_id == workspace_id
            )
        )
        return int(result.scalar() or 0)

    async def list_by_source(self, workspace_id: UUID, source_type: str, source_id: str) -> list[RuntimeAiLedger]:
        result = await self._db.execute(
            select(RuntimeAiLedger)
            .where(RuntimeAiLedger.workspace_id == workspace_id)
            .where(RuntimeAiLedger.source_type == source_type)
            .where(RuntimeAiLedger.source_id == source_id)
            .order_by(RuntimeAiLedger.created_at, RuntimeAiLedger.id)
        )
        return list(result.scalars().all())

    async def add_entry(
        self,
        *,
        workspace_id: UUID,
        entry_type: str,
        millicredits_delta: int,
        usage_product: Optional[str] = None,
        source_type: Optional[str] = None,
        source_id: Optional[str] = None,
        purchase_id: Optional[UUID] = None,
        project_id: Optional[UUID] = None,
        description: Optional[str] = None,
        metadata: Optional[dict] = None,
        allow_zero: bool = False,
    ) -> RuntimeAiLedger:
        if millicredits_delta == 0 and not allow_zero:
            raise ValueError("millicredits_delta must be non-zero")
        balance = await self.sum_balance(workspace_id) + millicredits_delta
        row = RuntimeAiLedger(
            id=uuid4(),
            workspace_id=workspace_id,
            entry_type=entry_type,
            millicredits_delta=millicredits_delta,
            balance_after=balance,
            usage_product=usage_product,
            source_type=source_type,
            source_id=source_id,
            purchase_id=purchase_id,
            project_id=project_id,
            description=description,
            metadata_json=metadata,
        )
        self._db.add(row)
        await self._db.flush()
        return row