import uuid
from typing import Literal

from sqlalchemy.ext.asyncio import AsyncSession

from app.crud import imp_inventory as inventory_crud
from app.crud import imp_tx as tx_crud
from app.crud import audit as audit_crud
from app.models.imp_batch import ImpBatch
from app.models.site import Site
from app.models.subject import Subject

TxType = Literal["RECEIPT", "DISPENSE", "RETURN", "RECONCILE", "DESTROY", "TRANSFER_IN", "TRANSFER_OUT"]

POSITIVE_TYPES = {"RECEIPT", "RETURN", "TRANSFER_IN"}
NEGATIVE_TYPES = {"DISPENSE", "DESTROY", "TRANSFER_OUT"}


async def _validate_entities(db: AsyncSession, study_id: uuid.UUID, site_id: uuid.UUID, batch_id: uuid.UUID, subject_id: uuid.UUID | None):
    site = await db.get(Site, site_id)
    if not site or site.study_id != study_id:
        raise ValueError("Site not found in study")
    batch = await db.get(ImpBatch, batch_id)
    if not batch or batch.study_id != study_id:
        raise ValueError("Batch not found in study")
    # batch status check could be extended
    if subject_id:
        subj = await db.get(Subject, subject_id)
        if not subj or subj.study_id != study_id:
            raise ValueError("Subject not found in study")


async def create_transaction_and_apply_inventory(
    db: AsyncSession,
    *,
    study_id: uuid.UUID,
    site_id: uuid.UUID,
    batch_id: uuid.UUID,
    subject_id: uuid.UUID | None,
    tx_type: TxType,
    quantity: int,
    tx_date,
    reference: str | None,
    notes: str | None,
    operator_id: uuid.UUID,
    operator_role: str,
) -> tuple:
    if quantity <= 0:
        raise ValueError("Quantity must be positive")
    await _validate_entities(db, study_id, site_id, batch_id, subject_id)

    async with db.begin_nested():
        inv = await inventory_crud.get_or_create_inventory(db, study_id, site_id, batch_id)
        delta = quantity if tx_type in POSITIVE_TYPES else -quantity
        new_qoh = inv.quantity_on_hand + delta
        if new_qoh < 0:
            raise ValueError("Insufficient inventory")

        inv.quantity_on_hand = new_qoh
        tx = await tx_crud.create_tx(
            db,
            study_id=study_id,
            site_id=site_id,
            batch_id=batch_id,
            subject_id=subject_id,
            tx_type=tx_type,
            quantity=quantity,
            tx_date=tx_date,
            reference=reference,
            notes=notes,
            created_by=operator_id,
        )
        await db.flush()
    detail = f"{tx_type} {quantity} units batch {batch_id} at site {site_id}, QOH -> {new_qoh}"
    await audit_crud.log_action(
        db,
        study_id=study_id,
        entity_type="imp_tx",
        entity_id=tx.id,
        action="CREATE_IMP_TX",
        detail=detail,
        operator_id=operator_id,
        operator_role=operator_role,
    )
    return tx