"""Aging Engine for AR/AP."""
from __future__ import annotations

from datetime import date
from decimal import Decimal
from uuid import UUID

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.receivable_payable import (
    AgingSnapshot,
    PayableInvoice,
    ReceivableInvoice,
)
from app.models.types import InvoiceStatus
from app.repositories.base import TenantBaseRepository


class AgingSnapshotRepository(TenantBaseRepository[AgingSnapshot]):
    model = AgingSnapshot


class AgingEngine:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session
        self.snapshot_repo = AgingSnapshotRepository(session)

    async def calculate_customer_aging(
        self, tenant_id: UUID, customer_id: UUID, as_of: date | None = None
    ) -> AgingSnapshot:
        as_of = as_of or date.today()
        stmt = select(ReceivableInvoice).where(
            ReceivableInvoice.tenant_id == tenant_id,
            ReceivableInvoice.customer_id == customer_id,
            ReceivableInvoice.status.in_([
                InvoiceStatus.ISSUED,
                InvoiceStatus.PARTIALLY_PAID,
            ]),
        )
        result = await self.session.execute(stmt)
        invoices = result.scalars().all()

        buckets = {
            "current": Decimal("0"),
            "1_30": Decimal("0"),
            "31_60": Decimal("0"),
            "61_90": Decimal("0"),
            "90_plus": Decimal("0"),
        }

        for inv in invoices:
            due = inv.due_date or inv.invoice_date
            days = (as_of - due).days
            amount = inv.remaining_amount
            if days <= 0:
                buckets["current"] += amount
            elif days <= 30:
                buckets["1_30"] += amount
            elif days <= 60:
                buckets["31_60"] += amount
            elif days <= 90:
                buckets["61_90"] += amount
            else:
                buckets["90_plus"] += amount

        total = sum(buckets.values())
        snapshot = AgingSnapshot(
            tenant_id=tenant_id,
            party_type="customer",
            party_id=customer_id,
            snapshot_date=as_of,
            current_amount=buckets["current"],
            days_1_30=buckets["1_30"],
            days_31_60=buckets["31_60"],
            days_61_90=buckets["61_90"],
            days_90_plus=buckets["90_plus"],
            total_outstanding=total,
        )
        return await self.snapshot_repo.add(snapshot)

    async def calculate_supplier_aging(
        self, tenant_id: UUID, supplier_id: UUID, as_of: date | None = None
    ) -> AgingSnapshot:
        as_of = as_of or date.today()
        stmt = select(PayableInvoice).where(
            PayableInvoice.tenant_id == tenant_id,
            PayableInvoice.supplier_id == supplier_id,
            PayableInvoice.status.in_([
                InvoiceStatus.ISSUED,
                InvoiceStatus.PARTIALLY_PAID,
            ]),
        )
        result = await self.session.execute(stmt)
        invoices = result.scalars().all()

        buckets = {
            "current": Decimal("0"),
            "1_30": Decimal("0"),
            "31_60": Decimal("0"),
            "61_90": Decimal("0"),
            "90_plus": Decimal("0"),
        }

        for inv in invoices:
            due = inv.due_date or inv.invoice_date
            days = (as_of - due).days
            amount = inv.remaining_amount
            if days <= 0:
                buckets["current"] += amount
            elif days <= 30:
                buckets["1_30"] += amount
            elif days <= 60:
                buckets["31_60"] += amount
            elif days <= 90:
                buckets["61_90"] += amount
            else:
                buckets["90_plus"] += amount

        total = sum(buckets.values())
        snapshot = AgingSnapshot(
            tenant_id=tenant_id,
            party_type="supplier",
            party_id=supplier_id,
            snapshot_date=as_of,
            current_amount=buckets["current"],
            days_1_30=buckets["1_30"],
            days_31_60=buckets["31_60"],
            days_61_90=buckets["61_90"],
            days_90_plus=buckets["90_plus"],
            total_outstanding=total,
        )
        return await self.snapshot_repo.add(snapshot)
