"""Unit tests for the sweep handler, using a fake S3. Run: python3 test_sweep.py

No AWS calls. Verifies the event path, the backstop path, idempotency on
duplicate delivery, URL decoding, and that we refuse to act on other buckets.
"""

import sys

from botocore.exceptions import ClientError

import lambda_function as lf

FAILURES = []


def check(name, actual, expected):
    if actual != expected:
        FAILURES.append("%s\n    expected: %r\n    actual:   %r" % (name, expected, actual))


class FakeS3(object):
    """Just enough S3 to exercise copy/delete/list."""

    def __init__(self, keys):
        self.buckets = {"bsd-wpe-logs": set(keys), "bsd-wpe-logs-private": set()}
        self.copies = []
        self.deletes = []

    def copy_object(self, Bucket, Key, CopySource):
        src, skey = CopySource["Bucket"], CopySource["Key"]
        if skey not in self.buckets[src]:
            raise ClientError({"Error": {"Code": "NoSuchKey", "Message": "nope"}}, "CopyObject")
        self.buckets[Bucket].add(Key)
        self.copies.append((skey, Bucket))

    def delete_object(self, Bucket, Key):
        self.buckets[Bucket].discard(Key)
        self.deletes.append(Key)

    def get_paginator(self, _name):
        outer = self

        class P(object):
            def paginate(self, Bucket, Prefix):
                items = sorted(k for k in outer.buckets[Bucket] if k.startswith(Prefix))
                yield {"Contents": [{"Key": k} for k in items]}
        return P()

    def list_objects_v2(self, Bucket, Prefix="", Delimiter=None):
        keys = [k for k in self.buckets[Bucket] if k.startswith(Prefix)]
        if not Delimiter:
            return {"KeyCount": len(keys)}
        prefixes, contents = set(), []
        for k in keys:
            rest = k[len(Prefix):]
            if Delimiter in rest:
                prefixes.add(Prefix + rest.split(Delimiter)[0] + Delimiter)
            else:
                contents.append({"Key": k})
        return {"CommonPrefixes": [{"Prefix": p} for p in sorted(prefixes)],
                "Contents": contents, "KeyCount": len(keys)}


def s3_event(key, bucket="bsd-wpe-logs"):
    return {"Records": [{"s3": {"bucket": {"name": bucket}, "object": {"key": key}}}]}


# --- event path moves exactly one object ---------------------------------
fake = FakeS3(["logs/nginx/a.log.gz", "logs/nginx/b.log.gz"])
lf.s3 = fake
r = lf.lambda_handler(s3_event("logs/nginx/a.log.gz"), None)
check("event moved 1", r["moved"], 1)
check("only a was copied", fake.copies, [("logs/nginx/a.log.gz", "bsd-wpe-logs-private")])
check("a gone from public", "logs/nginx/a.log.gz" in fake.buckets["bsd-wpe-logs"], False)
check("a now private", "logs/nginx/a.log.gz" in fake.buckets["bsd-wpe-logs-private"], True)
check("b untouched", "logs/nginx/b.log.gz" in fake.buckets["bsd-wpe-logs"], True)

# --- duplicate delivery is not an error ----------------------------------
r = lf.lambda_handler(s3_event("logs/nginx/a.log.gz"), None)
check("duplicate event -> gone, not raise", r, {"moved": 0, "skipped": 0, "gone": 1})

# --- URL-encoded keys -----------------------------------------------------
fake = FakeS3(["logs/nginx/my file+x.log.gz"])
lf.s3 = fake
r = lf.lambda_handler(s3_event("logs/nginx/my+file%2Bx.log.gz"), None)
check("url-decoded key moved", r["moved"], 1)
check("decoded correctly", fake.deletes, ["logs/nginx/my file+x.log.gz"])

# --- skip rules preserved -------------------------------------------------
fake = FakeS3(["logs/big.largefs", "logs/dir/"])
lf.s3 = fake
check("largefs skipped", lf.lambda_handler(s3_event("logs/big.largefs"), None)["skipped"], 1)
check("dir marker skipped", lf.lambda_handler(s3_event("logs/dir/"), None)["skipped"], 1)
check("nothing copied", fake.copies, [])

# --- .largefs is skipped wherever it sits, root included ------------------
# WP Engine keeps its validation files at the bucket ROOT and needs them there
# to keep delivering logs. They must never be moved or deleted, so pin the
# invariant at every depth rather than just under logs/.
for k in (".largefs", "wp-content.largefs", "logs/a.largefs", "logs/nginx/deep/b.largefs"):
    check("_skip(%r)" % k, lf._skip(k), True)

fake = FakeS3(["wp-content.largefs", "logs/nginx/a.log.gz"])
lf.s3 = fake
# event path: a root .largefs is skipped, not moved
r = lf.lambda_handler(s3_event("wp-content.largefs"), None)
check("root largefs -> skipped", r, {"moved": 0, "skipped": 1, "gone": 0})
check("root largefs still in public bucket", "wp-content.largefs" in fake.buckets["bsd-wpe-logs"], True)
check("root largefs never copied", fake.copies, [])
check("root largefs never deleted", fake.deletes, [])

# backstop path: root .largefs survives, a normal log still moves
r = lf.lambda_handler({}, None)
check("backstop still moves the normal log", r["moved"], 1)
check("root largefs survived backstop", "wp-content.largefs" in fake.buckets["bsd-wpe-logs"], True)
check("backstop deleted only the log", fake.deletes, ["logs/nginx/a.log.gz"])

# --- refuses events from the wrong bucket (no self-copy) ------------------
fake = FakeS3(["logs/nginx/a.log.gz"])
lf.s3 = fake
r = lf.lambda_handler(s3_event("logs/nginx/a.log.gz", bucket="bsd-wpe-logs-private"), None)
check("wrong bucket ignored", r, {"moved": 0, "skipped": 0, "gone": 0})
check("wrong bucket: no copy", fake.copies, [])

# --- backstop path (no Records) moves everything under logs/ --------------
fake = FakeS3(["logs/nginx/a.log.gz", "logs/nginx/b.log.gz", "other/c.txt"])
lf.s3 = fake
r = lf.lambda_handler({}, None)
check("backstop moved 2", r["moved"], 2)
check("backstop ignored non-logs prefix", "other/c.txt" in fake.buckets["bsd-wpe-logs"], True)

# --- backstop is a no-op in the healthy steady state ----------------------
r = lf.lambda_handler({}, None)
check("healthy steady state -> moved 0", r["moved"], 0)

# --- EventBridge scheduled event has no Records -> backstop ---------------
fake = FakeS3(["logs/nginx/z.log.gz"])
lf.s3 = fake
r = lf.lambda_handler({"source": "aws.events", "detail-type": "Scheduled Event"}, None)
check("eventbridge payload -> backstop", r["moved"], 1)

# --- a real failure propagates (so Lambda retries) ------------------------
class BoomS3(FakeS3):
    def copy_object(self, **kw):
        raise ClientError({"Error": {"Code": "InternalError", "Message": "boom"}}, "CopyObject")


lf.s3 = BoomS3(["logs/nginx/a.log.gz"])
try:
    lf.lambda_handler(s3_event("logs/nginx/a.log.gz"), None)
    check("real error raises", "did not raise", "RuntimeError")
except RuntimeError as exc:
    check("real error raises", "failed to move" in str(exc), True)


# --- WP Engine's other delivery prefix: wpe_logs/ -------------------------
# On 2026-07-17 delivery resumed under wpe_logs/ instead of logs/. The sweep
# watched only logs/, so ~331MB sat in the public bucket while it logged
# moved=0. Both prefixes are swept now, and both land under logs/ in DST so
# the parser (SOURCE_PREFIX=logs/nginx/) and its IAM policy need no change.
check("dest_key normalises wpe_logs/",
      lf.dest_key("wpe_logs/nginx/a.log.gz"), "logs/nginx/a.log.gz")
check("dest_key leaves logs/ alone",
      lf.dest_key("logs/nginx/a.log.gz"), "logs/nginx/a.log.gz")
check("dest_key leaves an unknown key alone", lf.dest_key("other/x"), "other/x")
check("only the prefix is rewritten, not the rest",
      lf.dest_key("wpe_logs/nginx/wpe_logs-in-name.log.gz"),
      "logs/nginx/wpe_logs-in-name.log.gz")

# event path
fake = FakeS3(["wpe_logs/nginx/a.log.gz"])
lf.s3 = fake
r = lf.lambda_handler(s3_event("wpe_logs/nginx/a.log.gz"), None)
check("wpe_logs event moved", r["moved"], 1)
check("landed under logs/ in DST", fake.copies, [("wpe_logs/nginx/a.log.gz", "bsd-wpe-logs-private")])
check("DST key is normalised", "logs/nginx/a.log.gz" in fake.buckets["bsd-wpe-logs-private"], True)
check("original gone from public", "wpe_logs/nginx/a.log.gz" in fake.buckets["bsd-wpe-logs"], False)

# backstop sweeps both prefixes, and still skips .largefs at the root
fake = FakeS3(["logs/nginx/a.log.gz", "wpe_logs/nginx/b.log.gz",
               "wpe_logs/error/c.log.gz", "254c9fdb.largefs"])
lf.s3 = fake
r = lf.lambda_handler({}, None)
check("backstop moves both prefixes", r["moved"], 3)
check("root .largefs untouched by the wider sweep",
      "254c9fdb.largefs" in fake.buckets["bsd-wpe-logs"], True)
check("everything landed under logs/",
      sorted(fake.buckets["bsd-wpe-logs-private"]),
      ["logs/error/c.log.gz", "logs/nginx/a.log.gz", "logs/nginx/b.log.gz"])

# a .largefs under a swept prefix is still skipped
fake = FakeS3(["wpe_logs/big.largefs"])
lf.s3 = fake
check("largefs skipped under wpe_logs/ too",
      lf.lambda_handler(s3_event("wpe_logs/big.largefs"), None)["skipped"], 1)
check("largefs not copied", fake.copies, [])

# --- the alarm that was missing -------------------------------------------
# moved=0 only means "nothing where I looked". This says whether we looked in
# the right place -- it is what would have caught the wpe_logs/ delivery.
fake = FakeS3(["logs/nginx/a.log.gz", "surprise_logs/nginx/x.log.gz",
               "surprise_logs/nginx/y.log.gz", "254c9fdb.largefs"])
lf.s3 = fake
check("unswept prefix is detected", lf.unswept_prefixes(), [("surprise_logs/", 2)])

fake = FakeS3(["logs/nginx/a.log.gz", "wpe_logs/nginx/b.log.gz", "254c9fdb.largefs"])
lf.s3 = fake
check("known prefixes raise no alarm", lf.unswept_prefixes(), [])
check("root .largefs is not an unswept prefix", lf.unswept_prefixes(), [])

if FAILURES:
    print("FAILED %d check(s):\n" % len(FAILURES))
    for f in FAILURES:
        print("  " + f)
    sys.exit(1)
print("all sweep checks passed")