"""Redact credentials, addresses, e-mail and home paths from agent trajectories.

Purpose: the stage between collecting trajectories and looking at them. An agent
    transcript is a recording of your machine talking about itself: the tool results
    contain whatever was in the files it read and whatever the commands it ran printed,
    which is where tokens, hostnames, addresses, e-mail and absolute paths live. This
    script walks every string in every episode, replaces what it recognises with a
    marker, drops the episodes whose leak cannot be safely redacted, and writes a report
    saying what it found and where. The pattern set is the one this course's own
    release check runs over its published files, extended with the two shapes that
    appear in transcripts rather than in source: bearer headers and connection strings.
Platform: all (standard library only; no model, no accelerator, no network)
Minimum memory: 8 GB nominally, and far less in practice: this is text in memory
Assumes: Python 3.10 or newer. The input is the JSON-lines episode file that
    collect-trajectories.py wrote. agentlog.py sits next to this file.

Usage: python3 scrub-trajectories.py --in raw/trajectories.jsonl \\
           --out clean/trajectories.jsonl --report scrub-report.json
       python3 scrub-trajectories.py --in raw/trajectories.jsonl --out clean/trajectories.jsonl \\
           --extra "client=\\bNorthwind\\b" --extra "internal-host=\\b[a-z]+\\.corp\\.invalid\\b"
       python3 scrub-trajectories.py --in clean/trajectories.jsonl --check
       python3 scrub-trajectories.py --list-patterns

A regular expression finds the things that have a shape. It does not find the things
that only you know are sensitive: a client's name, an unreleased product, the fact that
a particular repository exists. --extra is for those, and reading a sample of the output
is not optional. --check re-runs the pattern set over a file and exits non-zero if
anything still matches, which is the step to put in front of anything you publish.
"""
from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path
from typing import Any, Iterable

import agentlog

# Each entry is (id, compiled pattern, replacement, allow pattern or None). A match on a
# line the allow pattern also matches is left alone: those are the cases where the shape
# is present but the value is not private, such as the loopback address or a
# documentation e-mail domain.
BUILTIN: list[tuple[str, str, str, str | None]] = [
    ("private-key",
     r"-----BEGIN (?:RSA|OPENSSH|EC|DSA|PGP) PRIVATE KEY-----[\s\S]*?-----END [A-Z ]*PRIVATE KEY-----",
     "[redacted:private-key]", None),
    ("ssh-public-key",
     r"\b(?:ssh-(?:rsa|ed25519|dss)|ecdsa-sha2-nistp\d+) AAAA[0-9A-Za-z+/=]{20,}",
     "[redacted:ssh-public-key]", None),
    ("aws-key", r"\bAKIA[0-9A-Z]{16}\b", "[redacted:aws-key]", None),
    ("github-token", r"\bgh[pousr]_[A-Za-z0-9]{30,}\b", "[redacted:github-token]", None),
    ("slack-token", r"\bxox[abpr]-[A-Za-z0-9-]{10,}", "[redacted:slack-token]", None),
    ("api-secret", r"\bsk-(?:live|test|proj|ant)-[A-Za-z0-9_-]{16,}", "[redacted:api-secret]", None),
    ("huggingface-token", r"\bhf_[A-Za-z0-9]{30,}\b", "[redacted:huggingface-token]", None),
    ("ngc-key", r"\bnvapi-[A-Za-z0-9_-]{30,}\b", "[redacted:ngc-key]", None),
    ("bearer-header", r"(?i)\b(authorization\s*:\s*bearer)\s+\S+", r"\1 [redacted:bearer]", None),
    ("assigned-secret",
     r"(?i)\b(api[_-]?key|secret|token|passw(?:or)?d)(\s*[:=]\s*)(?:\"[^\"]{6,}\"|'[^']{6,}'|\S{6,})",
     r"\1\2[redacted:secret]",
     r"(?i)(?:placeholder|example|your[_-]|xxx|<[^>]+>|redacted)"),
    ("connection-string",
     r"\b[a-z][a-z0-9+.-]*://[^\s/@]+:[^\s/@]+@[^\s\"']+",
     "[redacted:connection-string]", None),
    # No exemption for documentation domains here, unlike the release check. A published
    # page may say example.com; a training set has no reason to carry any address at all,
    # and a rule with holes in it is a rule people stop reading.
    ("email", r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b", "[redacted:email]", None),
    ("ip-address",
     r"\b(?:(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)\.){3}(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)\b",
     "[redacted:address]", r"(?:127\.0\.0\.1|0\.0\.0\.0)"),
    ("mac-address", r"\b(?:[0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}\b", "[redacted:mac]", None),
    ("home-path",
     r"(?:/home/[a-z][a-z0-9_-]*|/Users/[A-Za-z][A-Za-z0-9_-]*|[A-Z]:\\Users\\[A-Za-z][A-Za-z0-9_-]*)",
     "[redacted:home]", None),
    ("root-login", r"\broot@[a-z0-9.-]+", "[redacted:root-login]", None),
]

# Leaks a marker cannot fix. A private key is not made safe by having its middle removed,
# because the fact that it was in this transcript is itself the finding, and the episode
# is worth nothing as training data anyway.
DEFAULT_DROP_ON = ("private-key",)


class Pattern:
    def __init__(self, name: str, regex: str, replacement: str, allow: str | None) -> None:
        self.name = name
        self.regex = re.compile(regex)
        self.replacement = replacement
        self.allow = re.compile(allow) if allow else None

    def apply(self, text: str) -> tuple[str, int, list[str]]:
        """Redact every match that the allow pattern does not exempt."""
        hits: list[str] = []

        def substitute(match: re.Match[str]) -> str:
            found = match.group(0)
            if self.allow and self.allow.search(found):
                return found
            hits.append(found[:12] + "…" if len(found) > 12 else found)
            return match.expand(self.replacement)

        return self.regex.sub(substitute, text), len(hits), hits


def build_patterns(extra: Iterable[str], redact_loopback: bool) -> list[Pattern]:
    patterns: list[Pattern] = []
    for name, regex, replacement, allow in BUILTIN:
        if name == "ip-address" and redact_loopback:
            allow = None
        patterns.append(Pattern(name, regex, replacement, allow))
    for item in extra:
        if "=" not in item:
            sys.exit(f"--extra needs NAME=REGEX, got {item!r}")
        name, regex = item.split("=", 1)
        try:
            patterns.append(Pattern(name.strip(), regex, f"[redacted:{name.strip()}]", None))
        except re.error as exc:
            sys.exit(f"--extra {name}: {exc}")
    return patterns


def scrub_value(value: Any, patterns: list[Pattern], counts: dict[str, int],
                samples: dict[str, list[str]], found: set[str]) -> Any:
    """Walk the whole object. Strings are redacted; everything else is rebuilt as it was."""
    if isinstance(value, str):
        text = value
        for pattern in patterns:
            text, n, hits = pattern.apply(text)
            if n:
                counts[pattern.name] = counts.get(pattern.name, 0) + n
                found.add(pattern.name)
                bucket = samples.setdefault(pattern.name, [])
                for hit in hits:
                    if len(bucket) < 5:
                        bucket.append(hit)
        return text
    if isinstance(value, list):
        return [scrub_value(v, patterns, counts, samples, found) for v in value]
    if isinstance(value, dict):
        return {k: scrub_value(v, patterns, counts, samples, found) for k, v in value.items()}
    return value


def read_episodes(path: Path) -> list[dict[str, Any]]:
    if not path.is_file():
        sys.exit(f"{path} does not exist; run collect-trajectories.py first")
    episodes = []
    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
        line = line.strip()
        if not line.startswith("{"):
            continue
        episodes.append(json.loads(line))
    if not episodes:
        sys.exit(f"{path} holds no episodes")
    return episodes


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--in", dest="input", default="raw/trajectories.jsonl")
    parser.add_argument("--out", default=None, help="where the scrubbed episodes go")
    parser.add_argument("--report", default=None, help="write the counts and samples here")
    parser.add_argument("--extra", action="append", default=[], metavar="NAME=REGEX",
                        help="one more pattern of your own; repeatable")
    parser.add_argument("--drop-on", action="append", default=None, metavar="NAME",
                        help=f"drop the whole episode on this pattern; default {DEFAULT_DROP_ON}")
    parser.add_argument("--redact-loopback", action="store_true",
                        help="also redact 127.0.0.1 and 0.0.0.0, which are kept by default")
    parser.add_argument("--check", action="store_true",
                        help="report what still matches and exit non-zero if anything does")
    parser.add_argument("--list-patterns", action="store_true")
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--notes", default=None)
    args = parser.parse_args()

    if args.list_patterns:
        print(f"{'name':<20} replacement")
        for name, _regex, replacement, allow in BUILTIN:
            note = "  (allow: " + allow + ")" if allow else ""
            print(f"{name:<20} {replacement}{note}")
        print("\nAdd your own with --extra NAME=REGEX. Client names, internal hostnames and "
              "project code names have no shape a regular expression can find on its own.")
        return

    patterns = build_patterns(args.extra, args.redact_loopback)
    drop_on = set(args.drop_on if args.drop_on is not None else DEFAULT_DROP_ON)
    episodes = read_episodes(Path(args.input))

    counts: dict[str, int] = {}
    samples: dict[str, list[str]] = {}
    kept: list[dict[str, Any]] = []
    dropped: list[dict[str, str]] = []

    for episode in episodes:
        found: set[str] = set()
        scrubbed = scrub_value(episode, patterns, counts, samples, found)
        blocking = sorted(found & drop_on)
        if blocking:
            dropped.append({"id": str(episode.get("id")), "patterns": ", ".join(blocking)})
            continue
        if found:
            scrubbed["scrubbed"] = sorted(found)
        kept.append(scrubbed)

    if args.check:
        total = sum(counts.values())
        if total:
            print(f"{total} match(es) still present in {args.input}:")
            for name in sorted(counts):
                print(f"  {name:<20} {counts[name]:>5}   e.g. {samples.get(name, [''])[0]}")
            sys.exit(1)
        print(f"{args.input}: nothing matched. That is a floor, not a proof: read a sample.")
        return

    if not args.out:
        parser.error("give --out, or --check to test a file that is already scrubbed")

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text("".join(json.dumps(e) + "\n" for e in kept), encoding="utf-8")

    report = {
        "input": args.input,
        "output": args.out,
        "episodes_in": len(episodes),
        "episodes_kept": len(kept),
        "episodes_dropped": dropped,
        "redactions_by_pattern": dict(sorted(counts.items())),
        "samples": {k: v for k, v in sorted(samples.items())},
        "extra_patterns": args.extra,
        "drop_on": sorted(drop_on),
        "loopback_redacted": args.redact_loopback,
    }
    if args.report:
        Path(args.report).write_text(json.dumps(report, indent=2), encoding="utf-8")

    print(f"episodes in:      {len(episodes)}")
    print(f"episodes kept:    {len(kept)}")
    print(f"episodes dropped: {len(dropped)}")
    if counts:
        print("\nredactions by pattern")
        for name in sorted(counts):
            print(f"  {name:<20} {counts[name]:>5}   e.g. {samples[name][0]}")
    else:
        print("\nNothing matched. Either the transcripts are clean or your agent's tools "
              "never printed anything private. Read ten of them and decide which.")
    for row in dropped:
        print(f"  dropped {row['id']}: {row['patterns']}")
    print(f"\nwritten to {out}")
    if args.report:
        print(f"report written to {args.report}")
    print("Now read a sample by hand. The patterns find shapes; they do not find the "
          "things only you know are sensitive.")

    if args.labbook:
        record = agentlog.record(
            labbook=args.labbook,
            lab="part-27/scrub-trajectories",
            model=None,
            dataset={"path": args.input, "sha256": agentlog.file_sha256(args.input),
                     "episodes": len(episodes)},
            data_lineage=agentlog.lineage(trajectories=args.input, scrub_report=args.report),
            hyperparameters={"extra_patterns": args.extra, "drop_on": sorted(drop_on),
                             "loopback_redacted": args.redact_loopback},
            seed=None, losses=None,
            scores={"episodes_kept": len(kept), "episodes_dropped": len(dropped),
                    "redactions": sum(counts.values())},
            config_path=__file__, notes=args.notes,
        )
        print(f"recorded run {record['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()