from decimal import Decimal
from sqlalchemy import select
from app.models.branch import Branch
from app.models.customer import Customer
from app.models.product import Product, ProductVariant
from app.models.receivable import Receivable, ReceivableCollection
from app.models.sale import Sale, SaleItem, SalePayment
from app.models.invoice import Invoice, InvoiceItem, InvoicePayment
from app.services.products import require_branch
from app.services.audit import record_audit

METHODS = {"CASH", "POS", "TRANSFER", "OTHER"}

def get_invoice(db, business_id, invoice_id, lock=False):
    q = select(Invoice).where(Invoice.id == invoice_id, Invoice.business_id == business_id)
    if lock: q = q.with_for_update()
    invoice = db.scalar(q)
    if not invoice: raise LookupError("Invoice not found.")
    return invoice

def _customer(db, business_id, customer_id):
    customer = db.get(Customer, customer_id)
    if not customer or customer.business_id != business_id or not customer.is_active: raise ValueError("Invalid customer.")
    return customer

def _item_rows(db, business_id, items):
    rows=[]
    for item in items:
        product = None
        if item.product_id:
            product = db.get(Product, item.product_id)
            if not product or product.business_id != business_id or not product.is_active: raise ValueError("Invalid product.")
        if item.variant_id:
            variant = db.get(ProductVariant, item.variant_id)
            if not variant or variant.business_id != business_id or (product and variant.product_id != product.id): raise ValueError("Invalid variant for product.")
        line = item.quantity * item.unit_price - item.discount
        if line < 0: raise ValueError("Item discount cannot exceed line price.")
        rows.append((item, line))
    return rows

def _set_totals(invoice, rows, discount, tax):
    subtotal=sum((line for _,line in rows), Decimal("0")); total=subtotal-discount+tax
    if total < 0: raise ValueError("Discount cannot exceed subtotal plus tax.")
    invoice.subtotal=subtotal; invoice.discount=discount; invoice.tax=tax; invoice.total=total; invoice.balance=total-invoice.amount_paid

def _sale_rows(db, sale):
    rows=[]
    for si in db.scalars(select(SaleItem).where(SaleItem.sale_id == sale.id)):
        product=db.get(Product, si.product_id)
        rows.append((type("Item", (), {"product_id":si.product_id,"variant_id":si.variant_id,"description":product.name if product else "Sale item","quantity":si.quantity,"unit_price":si.unit_price,"discount":si.discount})(), si.line_total))
    return rows

def create_invoice(db, business_id, user_id, data, commit=True):
    if data.branch_id: require_branch(db,business_id,data.branch_id)
    if data.sale_id:
        sale=db.scalar(select(Sale).where(Sale.id==data.sale_id,Sale.business_id==business_id).with_for_update())
        if not sale: raise ValueError("Invalid sale.")
        if not sale.customer_id: raise ValueError("Sale-linked invoice requires a sale customer.")
        if data.customer_id and data.customer_id != sale.customer_id: raise ValueError("Invoice customer must match sale customer.")
        if db.scalar(select(Invoice.id).where(Invoice.sale_id==sale.id)): raise ValueError("An invoice already exists for this sale.")
        customer=_customer(db,business_id,sale.customer_id); rows=_sale_rows(db,sale)
        invoice=Invoice(business_id=business_id,branch_id=sale.branch_id,customer_id=customer.id,sale_id=sale.id,invoice_number=data.invoice_number.strip(),status="DRAFT",issue_date=data.issue_date,due_date=data.due_date,subtotal=Decimal("0"),discount=Decimal("0"),tax=Decimal("0"),total=Decimal("0"),amount_paid=sale.amount_paid,balance=Decimal("0"),notes=data.notes,created_by=user_id)
        _set_totals(invoice,rows,Decimal("0"),Decimal("0"))
        if invoice.total != sale.total: raise ValueError("Sale totals are invalid.")
        db.add(invoice);db.flush()
    else:
        if not data.customer_id: raise ValueError("Customer is required.")
        customer=_customer(db,business_id,data.customer_id); rows=_item_rows(db,business_id,data.items)
        if not rows: raise ValueError("At least one invoice item is required.")
        invoice=Invoice(business_id=business_id,branch_id=data.branch_id,customer_id=customer.id,invoice_number=data.invoice_number.strip(),status="DRAFT",issue_date=data.issue_date,due_date=data.due_date,subtotal=Decimal("0"),discount=Decimal("0"),tax=Decimal("0"),total=Decimal("0"),amount_paid=Decimal("0"),balance=Decimal("0"),notes=data.notes,created_by=user_id)
        _set_totals(invoice,rows,data.discount,data.tax); db.add(invoice);db.flush()
        if invoice.balance: db.add(Receivable(business_id=business_id,customer_id=customer.id,invoice_id=invoice.id,original_amount=invoice.total,amount_paid=Decimal("0"),balance=invoice.total,due_date=invoice.due_date,status="OPEN"))
    for item,line in rows: db.add(InvoiceItem(invoice_id=invoice.id,product_id=item.product_id,variant_id=item.variant_id,description=item.description,quantity=item.quantity,unit_price=item.unit_price,discount=item.discount,line_total=line))
    if commit:
        record_audit(db,business_id,'INVOICE_CREATED','invoice',invoice.id,user_id,invoice.branch_id,{'invoice_number':invoice.invoice_number,'customer_id':str(invoice.customer_id),'total':str(invoice.total)}); db.commit(); db.refresh(invoice)
    return invoice

def update_invoice(db,business_id,invoice_id,data):
    invoice=get_invoice(db,business_id,invoice_id,True)
    if invoice.status != "DRAFT": raise ValueError("Only draft invoices can be edited.")
    values=data.model_dump(exclude_unset=True)
    if "branch_id" in values and values["branch_id"]: require_branch(db,business_id,values["branch_id"])
    if "customer_id" in values and values["customer_id"]: _customer(db,business_id,values["customer_id"])
    items=values.pop("items",None)
    for k,v in values.items(): setattr(invoice,k,v)
    if items is not None:
        rows=_item_rows(db,business_id,items)
        if not rows: raise ValueError("At least one invoice item is required.")
        for row in db.scalars(select(InvoiceItem).where(InvoiceItem.invoice_id==invoice.id)): db.delete(row)
        _set_totals(invoice,rows,invoice.discount,invoice.tax)
        for item,line in rows: db.add(InvoiceItem(invoice_id=invoice.id,product_id=item.product_id,variant_id=item.variant_id,description=item.description,quantity=item.quantity,unit_price=item.unit_price,discount=item.discount,line_total=line))
    if invoice.sale_id: raise ValueError("Sale-linked invoice financial items cannot be edited.")
    if items is None and ({"discount", "tax"} & values.keys()):
        rows=[(item, item.line_total) for item in db.scalars(select(InvoiceItem).where(InvoiceItem.invoice_id==invoice.id))]
        _set_totals(invoice, rows, invoice.discount, invoice.tax)
    receivable=db.scalar(select(Receivable).where(Receivable.invoice_id==invoice.id).with_for_update())
    if receivable:
        if invoice.customer_id != receivable.customer_id: receivable.customer_id=invoice.customer_id
        receivable.original_amount=invoice.total; receivable.balance=invoice.balance; receivable.due_date=invoice.due_date; receivable.status="PAID" if invoice.balance==0 else ("PARTIALLY_PAID" if invoice.amount_paid else "OPEN")
    db.commit();db.refresh(invoice);return invoice

def payment(db,business_id,invoice_id,user_id,data):
    if data.payment_method not in METHODS: raise ValueError("Invalid payment method.")
    invoice=get_invoice(db,business_id,invoice_id,True)
    if invoice.status == "CANCELLED": raise ValueError("Cancelled invoice cannot accept payment.")
    if invoice.balance <= 0 or data.amount > invoice.balance: raise ValueError("Payment exceeds invoice balance.")
    receivable = db.scalar(select(Receivable).where(Receivable.sale_id==invoice.sale_id if invoice.sale_id else Receivable.invoice_id==invoice.id).with_for_update())
    if receivable and data.amount > receivable.balance: raise ValueError("Payment exceeds receivable balance.")
    sale=None
    if invoice.sale_id:
        sale=db.scalar(select(Sale).where(Sale.id==invoice.sale_id,Sale.business_id==business_id).with_for_update())
        if not sale or sale.customer_id != invoice.customer_id: raise ValueError("Linked sale is invalid.")
    p=InvoicePayment(business_id=business_id,invoice_id=invoice.id,amount=data.amount,payment_method=data.payment_method,reference=data.reference,recorded_by=user_id);db.add(p)
    invoice.amount_paid += data.amount;invoice.balance -= data.amount;invoice.status="PAID" if invoice.balance==0 else "PARTIALLY_PAID"
    if receivable:
        receivable.amount_paid += data.amount;receivable.balance -= data.amount;receivable.status="PAID" if receivable.balance==0 else "PARTIALLY_PAID"
        db.add(ReceivableCollection(business_id=business_id,receivable_id=receivable.id,customer_id=invoice.customer_id,amount=data.amount,payment_method=data.payment_method,reference=data.reference,collected_by=user_id))
    if sale:
        sale.amount_paid += data.amount;sale.balance -= data.amount;sale.payment_status="PAID" if sale.balance==0 else "PARTIALLY_PAID"
        db.add(SalePayment(business_id=business_id,sale_id=sale.id,amount=data.amount,payment_method=data.payment_method,reference=data.reference,recorded_by=user_id))
    record_audit(db,business_id,'INVOICE_PAYMENT_RECORDED','invoice',invoice.id,user_id,invoice.branch_id,{'amount':str(data.amount),'payment_method':data.payment_method,'balance':str(invoice.balance)}); db.commit();db.refresh(p);return p

def cancel_invoice(db,business_id,invoice_id):
    invoice=get_invoice(db,business_id,invoice_id,True)
    if invoice.status == "CANCELLED": return invoice
    if invoice.amount_paid or invoice.status == "PAID": raise ValueError("Paid or financially active invoice cannot be cancelled.")
    if invoice.sale_id: raise ValueError("Sale-linked invoice cannot be cancelled without a sale reversal.")
    invoice.status="CANCELLED";record_audit(db,business_id,'INVOICE_CANCELLED','invoice',invoice.id,invoice.created_by,invoice.branch_id,{});db.commit();db.refresh(invoice);return invoice