#!/usr/bin/env python3
"""Estimate the unzipped size of a Terraform-packaged Lambda before deploying.

Mirrors what `deployment/terraform/modules/lambda_with_api_gateway/main.tf` does:
`pip install -r <requirements> -t <source_dir> ...` followed by zipping
`source_dir` minus `zip_excludes`. AWS Lambda's hard limit is 262144000 bytes
(250 MiB) unzipped. We install the dependencies into a scratch directory
instead of the real source tree so this can run locally or in CI without
mutating the repo.

Usage:
    python3 scripts/check_lambda_zip_size.py \
        --source-dir . \
        --requirements backend/app/requirements/requirements.txt \
        --exclude "**/tests/**" --exclude "**/.git/**" ...

Exits non-zero if the projected unzipped size exceeds --limit (or exceeds
--warn-pct of the limit, when --strict is not passed the warning is just
printed).
"""
from __future__ import annotations

import argparse
import fnmatch
import subprocess
import sys
import tempfile
from pathlib import Path

LAMBDA_UNZIPPED_LIMIT_BYTES = 262_144_000


def doublestar_match(pattern_parts: list[str], path_parts: list[str]) -> bool:
    if not pattern_parts:
        return not path_parts
    head, rest = pattern_parts[0], pattern_parts[1:]
    if head == "**":
        if doublestar_match(rest, path_parts):
            return True
        if path_parts and doublestar_match(pattern_parts, path_parts[1:]):
            return True
        return False
    if not path_parts:
        return False
    if fnmatch.fnmatchcase(path_parts[0], head):
        return doublestar_match(rest, path_parts[1:])
    return False


def is_excluded(rel_path: Path, excludes: list[str]) -> bool:
    parts = rel_path.as_posix().split("/")
    return any(doublestar_match(pattern.split("/"), parts) for pattern in excludes)


def sum_dir_size(root: Path, excludes: list[str] | None = None) -> int:
    total = 0
    for path in root.rglob("*"):
        if path.is_dir():
            continue
        rel = path.relative_to(root)
        if excludes and is_excluded(rel, excludes):
            continue
        total += path.stat().st_size
    return total


def install_dependencies(requirements: Path, target: Path) -> int:
    subprocess.run(
        [
            "pip",
            "install",
            "-r",
            str(requirements),
            "-t",
            str(target),
            "--platform",
            "manylinux2014_x86_64",
            "--implementation",
            "cp",
            "--python-version",
            "3.11",
            "--only-binary=:all:",
        ],
        check=True,
    )
    return sum_dir_size(target)


def human(n: int) -> str:
    return f"{n:,} bytes ({n / 1024 / 1024:.1f} MiB)"


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source-dir", type=Path, default=Path("."))
    parser.add_argument("--requirements", type=Path, default=None)
    parser.add_argument("--exclude", action="append", default=[], dest="excludes")
    parser.add_argument("--limit", type=int, default=LAMBDA_UNZIPPED_LIMIT_BYTES)
    parser.add_argument(
        "--warn-pct",
        type=float,
        default=85.0,
        help="Print a warning once projected size crosses this %% of the limit.",
    )
    parser.add_argument(
        "--strict",
        action="store_true",
        help="Exit non-zero on the warning threshold too, not just the hard limit.",
    )
    args = parser.parse_args()

    source_size = sum_dir_size(args.source_dir, args.excludes)

    deps_size = 0
    if args.requirements is not None:
        with tempfile.TemporaryDirectory(prefix="lambda-deps-") as tmp:
            deps_size = install_dependencies(args.requirements, Path(tmp))

    total = source_size + deps_size
    pct = total / args.limit * 100

    print(f"source files (post-exclude): {human(source_size)}")
    if args.requirements is not None:
        print(f"pip dependencies:            {human(deps_size)}")
    print(f"projected unzipped total:     {human(total)}  [{pct:.1f}% of limit]")
    print(f"AWS limit:                    {human(args.limit)}")

    if total > args.limit:
        print(
            f"\nFAIL: projected unzipped size exceeds the Lambda limit by "
            f"{human(total - args.limit)}.",
            file=sys.stderr,
        )
        return 1

    if pct >= args.warn_pct:
        msg = (
            f"\nWARNING: projected unzipped size is at {pct:.1f}% of the "
            f"{human(args.limit)} limit."
        )
        print(msg, file=sys.stderr)
        if args.strict:
            return 1

    return 0


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