"""
RAG tool Lambda  --  HNCB de-identification demo (steps 4-6).

The cloud agent calls this tool via AgentCore Gateway with a TOKEN, never a name.
This function:
  1. resolves the token against the on-prem vault  ->  real customer_id
  2. reads the raw customer record (stays here, never returned)
  3. returns a token-keyed, de-identified evidence package

If the token can't be resolved, it refuses -- it never guesses or leaks.
"""
import os
import boto3

ddb = boto3.resource("dynamodb")
VAULT = ddb.Table(os.environ["VAULT_TABLE"])
CUSTOMERS = ddb.Table(os.environ["CUSTOMERS_TABLE"])


def _tool_input(event):
    # AgentCore Gateway passes tool arguments in the event. Accept a few shapes
    # so this is robust across Gateway target styles during the demo.
    for key in ("input", "arguments", "parameters", "body"):
        if isinstance(event.get(key), dict):
            return event[key]
    return event


def lambda_handler(event, _context):
    args = _tool_input(event)
    token = args.get("customer_token")
    period_days = int(args.get("period_days", 90))

    if not token:
        return {"error": "customer_token is required"}

    # 1. resolve token -> customer_id (the reversible map lives only here)
    entry = VAULT.get_item(Key={"token": token}).get("Item")
    if not entry:
        return {"error": f"unknown token {token}; refusing to proceed"}
    customer_id = entry["value"] if entry.get("type") == "CUSTOMER" else entry.get("customer_id")
    if not customer_id:
        return {"error": "token did not resolve to a customer record"}

    # 2. read the raw record (NEVER returned to the caller)
    cust = CUSTOMERS.get_item(Key={"customer_id": customer_id}).get("Item")
    if not cust:
        return {"error": "customer record not found"}

    # 3. build a de-identified evidence package (summaries + categories only)
    return {
        "customer_token": token,
        "period": f"last {period_days} days",
        "activity_summary": cust.get("activity_summary", "n/a"),
        "amount_band": cust.get("amount_band", "n/a"),
        "product_type": cust.get("product_type", "n/a"),
        "risk_profile": cust.get("risk_profile", "n/a"),
        "sales_constraints": cust.get("sales_constraints", "n/a"),
        "note": "de-identified; no raw PII in this payload",
    }