#!/usr/bin/env python3
"""
anonymize.py — structure-preserving e-invoice anonymizer (UBL & CII)

Replaces commercially sensitive values (party names, tax/company IDs, bank
accounts, addresses, contacts, free-text notes) with realistic fakes while
preserving EVERYTHING validation depends on:

  • all amounts, quantities, percentages  -> untouched (arithmetic rules hold)
  • all code-list values                  -> untouched (codelist rules hold)
  • identifier FORMATS                    -> preserved (ABN gets a valid
    checksum, IBAN gets valid mod-97 check digits, VAT IDs keep country prefix)
  • referential consistency               -> the same original value maps to
    the same fake everywhere (PaymentID still matches the invoice number, etc.)
  • structured notes (#SKONTO#…, #ADU#…)  -> keywords kept, prose replaced

Result: the anonymized file produces the IDENTICAL validation outcome as the
original, so you can share it with a validator or consultant safely.

Usage:  python3 anonymize.py invoice.xml [-o out.xml] [--seed 42]
Runs 100% locally. No network. stdlib only.
"""

import argparse
import random
import re
import xml.etree.ElementTree as ET

# ---------------------------------------------------------------- fakes ----
COMPANIES = [
    "Acme Trading",
    "Bellbird Holdings",
    "Cormorant Supplies",
    "Dunmore Group",
    "Eastgate Industries",
    "Fernleaf Services",
    "Greywacke Partners",
    "Harbourline Co",
]
SUFFIX = {
    "AU": "Pty Ltd",
    "NZ": "Limited",
    "DE": "GmbH",
    "FR": "SARL",
    "GB": "Ltd",
    "NL": "B.V.",
    "BE": "BV",
    "default": "Ltd",
}
STREETS = [
    "14 Sample Street",
    "82 Placeholder Road",
    "7 Example Avenue",
    "230 Specimen Lane",
    "5 Redacted Terrace",
]
CITIES = ["Sampleton", "Mockville", "Testburg", "Placeholder Bay", "Redacton"]


class Anonymizer:
    def __init__(self, seed=None):
        self.rng = random.Random(seed)
        self.map = {}  # original value -> fake (consistency)
        self.company_i = 0

    def _memo(self, key, maker):
        if key not in self.map:
            self.map[key] = maker()
        return self.map[key]

    # -- identifier fakers that keep formats/checksums valid -----------------
    def fake_abn(self, orig):
        def make():
            while True:
                d = [self.rng.randint(1, 9)] + [self.rng.randint(0, 9) for _ in range(10)]
                w = [10, 1, 3, 5, 7, 9, 11, 13, 15, 17, 19]
                d0 = d[:]
                d0[0] -= 1
                if sum(a * b for a, b in zip(d0, w, strict=True)) % 89 == 0:
                    return "".join(map(str, d))

        return self._memo(("abn", orig), make)

    def fake_iban(self, orig):
        def make():
            cc = orig[:2] if re.match(r"[A-Z]{2}", orig) else "DE"
            bban = "".join(str(self.rng.randint(0, 9)) for _ in range(len(re.sub(r"\s", "", orig)) - 4 or 16))
            digits = "".join(str(int(c, 36)) for c in bban + cc + "00")
            check = 98 - int(digits) % 97
            return f"{cc}{check:02d}{bban}"

        return self._memo(("iban", orig), make)

    def fake_vat(self, orig):
        def make():
            m = re.match(r"([A-Z]{2})(.*)", orig.strip())
            prefix, body = (m.group(1), m.group(2)) if m else ("", orig)
            fake_body = "".join(
                str(self.rng.randint(0, 9))
                if c.isdigit()
                else (self.rng.choice("ABCDEFGHJKMNPQRSTUVWXYZ") if c.isalpha() else c)
                for c in body
            )
            return prefix + fake_body

        return self._memo(("vat", orig), make)

    def fake_digits_like(self, orig):
        return self._memo(
            ("num", orig), lambda: "".join(str(self.rng.randint(0, 9)) if c.isdigit() else c for c in orig)
        )

    def fake_company(self, orig, country="default"):
        def make():
            base = COMPANIES[self.company_i % len(COMPANIES)]
            self.company_i += 1
            return f"{base} {SUFFIX.get(country, SUFFIX['default'])}"

        return self._memo(("co", orig), make)

    def fake_email(self, orig):
        return self._memo(("mail", orig), lambda: f"contact{self.rng.randint(100, 999)}@anon-example.com")

    def fake_ref(self, orig):
        return self._memo(
            ("ref", orig), lambda: "ANON-" + "".join(self.rng.choice("0123456789") for _ in range(6))
        )

    def fake_text(self, orig):
        """Free-text notes: keep structured #KEYWORD#k=v# tokens, redact prose."""

        def repl(seg):
            return seg if seg.startswith("#") else "Redacted free-text note."

        parts = re.split(r"(#[A-Z]+#[^#]*(?:#[^#]*)*#?)", orig)
        if any(p.startswith("#") for p in parts):
            return "".join(p if p.startswith("#") else ("" if not p.strip() else " ") for p in parts) or orig
        return "Redacted free-text note."


# ------------------------------------------------------------- traversal ----
UBL_SENSITIVE_LEAVES = {
    "Name",
    "RegistrationName",
    "StreetName",
    "AdditionalStreetName",
    "CityName",
    "PostalZone",
    "CountrySubentity",
    "ElectronicMail",
    "Telephone",
    "Telefax",
    "Line",
    "Department",
    "FamilyName",
    "FirstName",
    "JobTitle",
}
CII_NAME_LEAVES = {"Name", "LineOne", "LineTwo", "CityName", "PostcodeCode", "PersonName", "DepartmentName"}
ID_LEAVES = {
    "CompanyID",
    "EndpointID",
    "IBANID",
    "ProprietaryID",
    "URIID",
    "AccountID",
    "PayerPartyDebtorFinancialAccountID",
}


def localname(tag):
    return tag.rsplit("}", 1)[-1]


def guess_country(root):
    for el in root.iter():
        if localname(el.tag) == "IdentificationCode" and el.text and len(el.text.strip()) == 2:
            return el.text.strip()
    return "default"


def anonymize_tree(root, an):
    country = guess_country(root)
    invoice_id = None
    # pass 1: find document ID (first top-level cbc:ID / ram ExchangedDocument ID)
    for el in root:
        if localname(el.tag) == "ID" and el.text:
            invoice_id = el.text.strip()
            break
    if invoice_id is None:
        for el in root.iter():
            if localname(el.tag) == "ExchangedDocument":
                for c in el:
                    if localname(c.tag) == "ID" and c.text:
                        invoice_id = c.text.strip()
                break
    stack = [(root, [])]
    while stack:
        el, path = stack.pop()
        ln = localname(el.tag)
        text = (el.text or "").strip()
        parent = path[-1] if path else ""

        if text:
            if ln in ("Name", "RegistrationName") and parent in (
                "PartyName",
                "PartyLegalEntity",
                "Party",
                "SellerTradeParty",
                "BuyerTradeParty",
                "PayeeTradeParty",
                "PayeeFinancialAccount",
                "PayeePartyCreditorFinancialAccount",
                "Contact",
                "DefinedTradeContact",
            ):
                el.text = (
                    an.fake_company(text, country)
                    if parent != "Contact" and parent != "DefinedTradeContact"
                    else an._memo(("person", text), lambda: "Alex Sample")
                )
            elif ln in UBL_SENSITIVE_LEAVES and parent in (
                "PostalAddress",
                "PostalTradeAddress",
                "Contact",
                "DefinedTradeContact",
                "Person",
                "AddressLine",
                "RegistrationAddress",
            ):
                if ln in ("StreetName", "AdditionalStreetName", "Line", "LineOne", "LineTwo"):
                    el.text = an._memo(("street", text), lambda: an.rng.choice(STREETS))
                elif ln in ("CityName",):
                    el.text = an._memo(("city", text), lambda: an.rng.choice(CITIES))
                elif ln in ("PostalZone", "PostcodeCode"):
                    el.text = an.fake_digits_like(text)
                elif ln == "ElectronicMail":
                    el.text = an.fake_email(text)
                elif ln in ("Telephone", "Telefax"):
                    el.text = an.fake_digits_like(text)
                elif ln in ("Name", "FamilyName", "FirstName", "PersonName"):
                    el.text = an._memo(("person", text), lambda: "Alex Sample")
            elif ln in CII_NAME_LEAVES and parent in (
                "SellerTradeParty",
                "BuyerTradeParty",
                "PostalTradeAddress",
                "DefinedTradeContact",
                "SpecifiedLegalOrganization",
            ):
                if ln == "Name":
                    el.text = an.fake_company(text, country)
                elif ln in ("LineOne", "LineTwo"):
                    el.text = an._memo(("street", text), lambda: an.rng.choice(STREETS))
                elif ln == "CityName":
                    el.text = an._memo(("city", text), lambda: an.rng.choice(CITIES))
                elif ln == "PostcodeCode":
                    el.text = an.fake_digits_like(text)
                else:
                    el.text = an._memo(("person", text), lambda: "Alex Sample")
            elif ln in ID_LEAVES or (
                ln == "ID"
                and parent
                in (
                    "PayeeFinancialAccount",
                    "PayeePartyCreditorFinancialAccount",
                    "FinancialInstitutionBranch",
                    "PayerPartyDebtorFinancialAccount",
                )
            ):
                scheme = el.get("schemeID", "")
                clean = re.sub(r"\s", "", text)
                if scheme == "0151" or (country == "AU" and re.fullmatch(r"\d{11}", clean)):
                    el.text = an.fake_abn(clean)
                elif re.fullmatch(r"[A-Z]{2}\d{2}[A-Z0-9]{10,30}", clean):
                    el.text = an.fake_iban(clean)
                elif re.match(r"[A-Z]{2}", clean) and any(c.isdigit() for c in clean):
                    el.text = an.fake_vat(clean)
                else:
                    el.text = an.fake_digits_like(text)
            elif ln == "CompanyID":
                el.text = an.fake_vat(text) if re.match(r"[A-Z]{2}", text) else an.fake_digits_like(text)
            elif (
                ln in ("ID", "PaymentID", "BuyerReference", "OrderReference")
                and invoice_id
                and text == invoice_id
            ):
                el.text = an.fake_ref(text)  # consistent doc-number mapping
            elif ln == "Note":
                el.text = an.fake_text(text)

        for child in el:
            stack.append((child, path + [ln]))


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[1])
    ap.add_argument("input")
    ap.add_argument("-o", "--output")
    ap.add_argument("--seed", type=int, help="deterministic output for repeatable tests")
    a = ap.parse_args()
    # preserve namespaces on output
    events = ET.iterparse(a.input, events=["start-ns"])
    namespaces = []
    for _, (prefix, uri) in events:
        namespaces.append((prefix, uri))
        ET.register_namespace(prefix, uri)
    tree = ET.parse(a.input)
    anonymize_tree(tree.getroot(), Anonymizer(a.seed))
    out = a.output or re.sub(r"\.xml$", "", a.input) + ".anon.xml"
    tree.write(out, encoding="UTF-8", xml_declaration=True)
    print(f"anonymized -> {out}")
    print("Verify: run your validator on both files — the results should be identical.")


if __name__ == "__main__":
    main()