models.py

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List, Literal, Optional

# ── Line item ────────────────────────────────────────────────────────────────

@dataclass
class InvoiceLineItem:
    product_raw:        str
    product_normalized: Optional[str]  = None
    product_id:         Optional[str]  = None
    quantity:           float          = 0.0
    unit_price:         float          = 0.0
    gst_rate:           float          = 0.0
    line_total:         float          = 0.0


# ── Invoice JSON ──────────────────────────────────────────────────────────────

@dataclass
class InvoiceJSON:
    invoice_number:      Optional[str]             = None
    supplier:            Optional[str]             = None
    date:                Optional[str]             = None
    items:               List[InvoiceLineItem]     = field(default_factory=list)
    grand_total:         float                     = 0.0
    extraction_warnings: List[str]                 = field(default_factory=list)
    unmatched_products:  List[str]                 = field(default_factory=list)


NormalisedInvoice = InvoiceJSON


# ── Anomaly flags ─────────────────────────────────────────────────────────────

FlagType = Literal[
    "price_anomaly",
    "delivery_shortage",
    "duplicate_charge",
    "gst_mismatch",
    "duplicate_invoice",
]


@dataclass
class AnomalyFlag:
    flag_type:    FlagType
    product_id:   Optional[str]  = None
    product_name: Optional[str]  = None
    amount_inr:   float          = 0.0
    description:  str            = ""
    action_item:  str            = ""
    metadata:     dict           = field(default_factory=dict)


# ── Leakage report ────────────────────────────────────────────────────────────

@dataclass
class LeakageReport:
    audit_run_id:          str
    invoice_number:        Optional[str]       = None
    supplier:              Optional[str]       = None
    date:                  Optional[str]       = None
    anomaly_flags:         List[AnomalyFlag]  = field(default_factory=list)
    total_leakage_inr:     float              = 0.0
    action_items:          List[str]          = field(default_factory=list)
    has_delivery_data:     bool               = False
    unexpected_deliveries: List[str]          = field(default_factory=list)
    low_confidence_photos: List[str]          = field(default_factory=list)
    skipped_price_checks:  int               = 0
    skipped_gst_checks:    int               = 0
    passed_all_checks:     bool              = False


# ── Agent trace ───────────────────────────────────────────────────────────────

@dataclass
class AgentTraceEntry:
    agent_name:      str
    agent_version:   str
    audit_run_id:    str
    timestamp_start: str            # ISO 8601 UTC
    timestamp_end:   str
    duration_ms:     int
    input_summary:   str
    output_summary:  str


CANONICAL_AGENT_ORDER: List[str] = [
    "Invoice_Extractor",
    "Product_Matcher",
    "Pricing_Agent",
    "Visual_Counter",
    "Reconciliation_Agent",
    "Savings_Agent",
]


# ── Catalog entry ─────────────────────────────────────────────────────────────

@dataclass
class CatalogEntry:
    product_id:     str
    canonical_name: str
    hsn_code:       str
    gst_rate:       float
    category:       str
    brand:          str
    common_aliases: List[str] = field(default_factory=list)


# ── Type aliases ──────────────────────────────────────────────────────────────

DeliveryCountMap = Dict[str, int]


# ── Progress streaming ────────────────────────────────────────────────────────

AuditStage = Literal[
    "extracting",
    "normalising",
    "checking_prices",
    "counting_products",
    "reconciling",
    "generating_report",
    "complete",
    "error",
]


@dataclass
class AuditProgressUpdate:
    stage:      AuditStage
    message:    str        = ""
    agent_name: str        = ""
    duration_ms: int       = 0


# ── Pricing agent metadata ────────────────────────────────────────────────────

@dataclass
class PricingAgentMeta:
    skipped_price_checks: int = 0
    skipped_gst_checks:   int = 0