procurement-ingest/lambdas/site_extractor/handler.py

"""
DynamoDB Streams processor that extracts Amazon site codes and addresses
from purchase order records and upserts them into the verified-sites table.
"""

import logging
import os
import re
from datetime import datetime, timezone

import boto3
from boto3.dynamodb.types import TypeDeserializer

logger = logging.getLogger()
logger.setLevel(logging.INFO)

dynamodb = boto3.client("dynamodb")
deserializer = TypeDeserializer()

VERIFIED_SITES_TABLE = os.environ.get("VERIFIED_SITES_TABLE", "verified-sites")

SITE_CODE_PATTERN = re.compile(r"[A-Z]{2,4}\d{1,2}")


def deserialize_image(image: dict) -> dict:
    return {k: deserializer.deserialize(v) for k, v in image.items()}


def extract_site_code(record: dict) -> str | None:
    # Prefer the explicit site_code field (set by updated extraction prompt)
    direct = (record.get("site_code") or "").strip().upper()
    if direct and SITE_CODE_PATTERN.match(direct):
        return direct

    ship_to = record.get("ship_to") or {}
    ship_to_name = ship_to.get("name", "")

    if ship_to_name:
        m = re.search(r"\(([A-Z0-9]{3,5})\)", ship_to_name)
        if m and SITE_CODE_PATTERN.match(m.group(1)):
            return m.group(1)

        m = re.search(r"(?:LLC|Inc)\s*-\s*([A-Z0-9]{3,5})\b", ship_to_name)
        if m and SITE_CODE_PATTERN.match(m.group(1)):
            return m.group(1)

        m = re.search(
            r"(?:Station|DS)\s*[-–]?\s*([A-Z0-9]{3,5})\b", ship_to_name
        )
        if m and SITE_CODE_PATTERN.match(m.group(1)):
            return m.group(1)

        m = re.match(r"^([A-Z0-9]{3,5})\s*-\s*Amazon", ship_to_name)
        if m and SITE_CODE_PATTERN.match(m.group(1)):
            return m.group(1)

    for item in record.get("line_items") or []:
        desc = item.get("description", "")
        m = re.match(r"^([A-Z]{1,4}[0-9]{1,2}|[A-Z]{4,5})\b", desc)
        if m and SITE_CODE_PATTERN.match(m.group(1)):
            return m.group(1)

    return None


def parse_address(record: dict) -> dict:
    result = {
        "address": None,
        "city": None,
        "state": None,
        "zip": None,
        "fullAddress": None,
    }

    ship_to = record.get("ship_to") or {}

    # Prefer structured fields if the extraction prompt provided them
    if ship_to.get("street") and ship_to.get("state"):
        result["address"] = ship_to["street"]
        result["city"] = ship_to.get("city")
        result["state"] = ship_to["state"]
        result["zip"] = ship_to.get("zip")
        parts = [p for p in [result["address"], result["city"], result["state"], result["zip"]] if p]
        result["fullAddress"] = ", ".join(parts)
        return result

    address_text = ship_to.get("address", "")
    if not address_text:
        return result

    clean = re.sub(r",?\s*United States\s*$", "", address_text.strip())
    if not clean:
        return result

    # Match "city, ST ZIP" at the end of the string
    m = re.search(r",\s*([^,]+?),\s+([A-Z]{2})\s+(\d{5}(?:-\d{4})?)\s*$", clean)
    if m:
        result["address"] = clean[: m.start()].strip()
        result["city"] = m.group(1).strip()
        result["state"] = m.group(2)
        result["zip"] = m.group(3)
        result["fullAddress"] = f"{result['address']}, {result['city']}, {result['state']} {result['zip']}"
        return result

    # Fallback: try newline-separated format
    lines = [l.strip() for l in clean.split("\n") if l.strip()]
    city_re = re.compile(r"^(.+?),\s+([A-Z]{2})\s+(\d{5}(?:-\d{4})?)")
    for i, line in enumerate(lines):
        m = city_re.match(line)
        if m:
            result["address"] = ", ".join(lines[:i]) if i > 0 else None
            result["city"] = m.group(1)
            result["state"] = m.group(2)
            result["zip"] = m.group(3)
            street = result["address"] or result["city"]
            result["fullAddress"] = f"{street}, {result['city']}, {result['state']} {result['zip']}"
            return result

    result["address"] = clean
    return result


def upsert_site(
    site_code: str,
    address: dict,
    po_number: str,
    location_code: str | None,
):
    now = datetime.now(timezone.utc).isoformat()

    expr_names = {}
    expr_values = {
        ":now": {"S": now},
        ":one": {"N": "1"},
        ":po_set": {"SS": [po_number]},
    }

    set_parts = ["lastSeenAt = :now"]
    if address.get("address"):
        set_parts.append("#addr = :addr")
        expr_names["#addr"] = "address"
        expr_values[":addr"] = {"S": address["address"]}
    if address.get("city"):
        set_parts.append("city = :city")
        expr_values[":city"] = {"S": address["city"]}
    if address.get("state"):
        set_parts.append("#st = :state_val")
        expr_names["#st"] = "state"
        expr_values[":state_val"] = {"S": address["state"]}
    if address.get("zip"):
        set_parts.append("zip = :zip")
        expr_values[":zip"] = {"S": address["zip"]}
    if address.get("fullAddress"):
        set_parts.append("fullAddress = :full")
        expr_values[":full"] = {"S": address["fullAddress"]}
    if location_code:
        set_parts.append("locationCode = :loc")
        expr_values[":loc"] = {"S": location_code}

    update_expr = f"SET {', '.join(set_parts)} ADD poCount :one, sourcePOs :po_set"

    kwargs = {
        "TableName": VERIFIED_SITES_TABLE,
        "Key": {"siteCode": {"S": site_code}},
        "UpdateExpression": update_expr,
        "ExpressionAttributeValues": expr_values,
    }
    if expr_names:
        kwargs["ExpressionAttributeNames"] = expr_names

    dynamodb.update_item(**kwargs)


def handler(event, context):
    processed = 0
    extracted = 0
    skipped = 0

    for record in event.get("Records", []):
        if record["eventName"] not in ("INSERT", "MODIFY"):
            continue

        new_image = record.get("dynamodb", {}).get("NewImage")
        if not new_image:
            continue

        processed += 1
        po = deserialize_image(new_image)
        po_number = po.get("po_number", "unknown")

        site_code = extract_site_code(po)
        if not site_code:
            skipped += 1
            logger.info("No site code for PO %s", po_number)
            continue

        address = parse_address(po)
        location_code = (po.get("ship_to") or {}).get("location_code")

        upsert_site(site_code, address, po_number, location_code)
        extracted += 1
        logger.info("Upserted site %s from PO %s", site_code, po_number)

    logger.info(
        "Batch complete: %d processed, %d extracted, %d skipped",
        processed,
        extracted,
        skipped,
    )