#!/usr/bin/env python3

from __future__ import annotations

import argparse
import json
from pathlib import Path
import re
from typing import Any

import yaml


SENSITIVE_KEY_RE = re.compile(
    r"(?i)(authorization|cookie|set-cookie|token|api[-_]?key|secret|password|session(_id)?|chatgpt[-_]account[-_]id)"
)
BEARER_RE = re.compile(r"(?i)\bBearer\s+[A-Za-z0-9._-]+\b")
JWT_RE = re.compile(r"\beyJ[A-Za-z0-9_-]+\.[A-Za-z0-9._-]+\.[A-Za-z0-9._-]+\b")
EMAIL_RE = re.compile(r"\b[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}\b", re.IGNORECASE)
HOME_PATH_RE = re.compile(r"/Users/[^/\s]+")
REFRESH_TOKEN_RE = re.compile(r'("refresh_token"\s*:\s*")[^"]+(")', re.IGNORECASE)
ACCESS_TOKEN_RE = re.compile(r'("access_token"\s*:\s*")[^"]+(")', re.IGNORECASE)
GENERIC_TOKEN_RE = re.compile(
    r'((?:access_token|refresh_token|id_token|api_key|session_id|token)=)([^&\s]+)',
    re.IGNORECASE,
)

REPLACEMENTS = {
    "secret": "[REDACTED]",
    "email": "[REDACTED_EMAIL]",
    "path": "/Users/[REDACTED_USER]",
}


def scrub_string(value: str) -> str:
    value = BEARER_RE.sub("Bearer [REDACTED]", value)
    value = JWT_RE.sub("[REDACTED_JWT]", value)
    value = EMAIL_RE.sub(REPLACEMENTS["email"], value)
    value = HOME_PATH_RE.sub(REPLACEMENTS["path"], value)
    value = REFRESH_TOKEN_RE.sub(r'\1[REDACTED]\2', value)
    value = ACCESS_TOKEN_RE.sub(r'\1[REDACTED]\2', value)
    value = GENERIC_TOKEN_RE.sub(r"\1[REDACTED]", value)
    return value


def scrub_data(value: Any, key: str | None = None) -> Any:
    if isinstance(value, dict):
        scrubbed: dict[str, Any] = {}
        for child_key, child_value in value.items():
            if SENSITIVE_KEY_RE.search(str(child_key)):
                if isinstance(child_value, str):
                    scrubbed[child_key] = REPLACEMENTS["secret"]
                else:
                    scrubbed[child_key] = scrub_data(child_value, child_key)
                continue
            scrubbed[child_key] = scrub_data(child_value, str(child_key))
        return scrubbed
    if isinstance(value, list):
        return [scrub_data(item, key) for item in value]
    if isinstance(value, str):
        if key and SENSITIVE_KEY_RE.search(key):
            return REPLACEMENTS["secret"]
        return scrub_string(value)
    return value


def scrub_json_like_text(text: str) -> str:
    try:
        payload = json.loads(text)
    except json.JSONDecodeError:
        return scrub_string(text)
    return json.dumps(scrub_data(payload), indent=2) + "\n"


def scrub_jsonl_text(text: str) -> str:
    lines: list[str] = []
    for raw_line in text.splitlines():
        if not raw_line.strip():
            continue
        try:
            payload = json.loads(raw_line)
            lines.append(json.dumps(scrub_data(payload), separators=(",", ":")))
        except json.JSONDecodeError:
            lines.append(scrub_string(raw_line))
    return "\n".join(lines) + ("\n" if lines else "")


def scrub_yaml_text(text: str) -> str:
    try:
        payload = yaml.safe_load(text)
    except yaml.YAMLError:
        return scrub_string(text)
    return yaml.safe_dump(scrub_data(payload), sort_keys=False)


def scrub_file(path: Path) -> bool:
    original = path.read_text(encoding="utf-8")
    suffixes = path.suffixes
    if path.name.endswith(".jsonl"):
        scrubbed = scrub_jsonl_text(original)
    elif ".json" in suffixes:
        scrubbed = scrub_json_like_text(original)
    elif any(suffix in suffixes for suffix in (".yaml", ".yml")):
        scrubbed = scrub_yaml_text(original)
    else:
        scrubbed = scrub_string(original)
    if scrubbed == original:
        return False
    path.write_text(scrubbed, encoding="utf-8")
    return True


def iter_target_files(root: Path) -> list[Path]:
    return sorted(path for path in root.rglob("*") if path.is_file())


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="rho_fixture_scrub.py")
    parser.add_argument("paths", nargs="+", type=Path, help="Files or directories to scrub in place")
    return parser


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    changed = 0
    for target in args.paths:
        if target.is_dir():
            candidates = iter_target_files(target)
        elif target.is_file():
            candidates = [target]
        else:
            raise SystemExit(f"path not found: {target}")
        for candidate in candidates:
            if scrub_file(candidate):
                changed += 1
                print(candidate)
    print(f"scrubbed files: {changed}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
