#!/usr/bin/env python3
"""Estimate the unzipped size of the fastapi Lambda before deploying.

Mirrors what `deployment/terraform/modules/lambda_with_api_gateway/main.tf` does
for the `ara_fast_api` lambda (`deployment/terraform/lambda/fast-api`):
`pip install -r requirements.txt -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 (from the repo root):
    python3 backend/app/requirements/check_lambda_zip_size.py --source-dir .

--requirements defaults to the requirements.txt next to this script, and
--exclude defaults to whatever `zip_excludes` is set to in
`deployment/terraform/modules/lambda_with_api_gateway/variables.tf` — read
straight from that file, so this can never silently drift from what
Terraform actually excludes. Pass --exclude explicitly to override.

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 re
import subprocess
import sys
import tempfile
from pathlib import Path

LAMBDA_UNZIPPED_LIMIT_BYTES = 262_144_000

REPO_ROOT = Path(__file__).resolve().parents[3]
DEFAULT_REQUIREMENTS = Path(__file__).resolve().parent / "requirements.txt"
DEFAULT_TF_VARIABLES_FILE = (
    REPO_ROOT / "deployment/terraform/modules/lambda_with_api_gateway/variables.tf"
)


def parse_zip_excludes_from_tf(tf_variables_file: Path) -> list[str]:
    text = tf_variables_file.read_text()
    match = re.search(
        r'variable\s+"zip_excludes"\s*{.*?default\s*=\s*\[(.*?)\]', text, re.S
    )
    if not match:
        raise ValueError(
            f'could not find variable "zip_excludes" default in {tf_variables_file}'
        )
    return re.findall(r'"([^"]+)"', match.group(1))


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=REPO_ROOT)
    parser.add_argument("--requirements", type=Path, default=DEFAULT_REQUIREMENTS)
    parser.add_argument(
        "--skip-deps",
        action="store_true",
        help="Skip the pip install step and only measure source files.",
    )
    parser.add_argument(
        "--exclude",
        action="append",
        default=None,
        dest="excludes",
        help="Repeatable. Defaults to zip_excludes read from --tf-variables-file.",
    )
    parser.add_argument("--tf-variables-file", type=Path, default=DEFAULT_TF_VARIABLES_FILE)
    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()

    excludes = args.excludes
    if excludes is None:
        excludes = parse_zip_excludes_from_tf(args.tf_variables_file)
        print(f"excludes (from {args.tf_variables_file.relative_to(REPO_ROOT)}): {excludes}")

    source_size = sum_dir_size(args.source_dir, excludes)

    deps_size = 0
    if not args.skip_deps:
        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 not args.skip_deps:
        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())