#!/usr/bin/env bash

# This script updates the sub dependencies in requirements.txt based on the base image specified in Dockerfile.

set -euo pipefail

req_file="requirements.txt"
dockerfile="Dockerfile"

if [[ ! -f "$req_file" ]]; then
  echo "Missing $req_file in $(pwd)" >&2
  exit 1
fi

if [[ ! -f "$dockerfile" ]]; then
  echo "Missing $dockerfile in $(pwd)" >&2
  exit 1
fi

image=$(awk 'toupper($1)=="FROM" {print $2; exit}' "$dockerfile")
if [[ -z "${image}" ]]; then
  echo "Could not determine base image from $dockerfile" >&2
  exit 1
fi

tmp_req_path=$(mktemp "$PWD/.tmp_requirements.XXXXXX")
freeze_path=$(mktemp "$PWD/.tmp_freeze.XXXXXX")
cleanup() {
  rm -f "$tmp_req_path" "$freeze_path"
}
trap cleanup EXIT

python3 - "$req_file" "$tmp_req_path" <<'PY'
import re
import sys

req_path = sys.argv[1]
out_path = sys.argv[2]

with open(req_path, "r", encoding="utf-8") as f:
    lines = f.readlines()

out_lines = []
inside_sub = False
for line in lines:
    stripped = line.strip()
    lower = stripped.lower()
    if lower.startswith("#") and ("sub depend" in lower or "sub-depend" in lower):
        inside_sub = True
        out_lines.append(line)
        continue
    if lower.startswith("#") and inside_sub:
        inside_sub = False
        out_lines.append(line)
        continue
    if inside_sub:
        if stripped and not stripped.startswith("#"):
            out_lines.append("# " + line.lstrip("# "))
        else:
            out_lines.append(line)
    else:
        out_lines.append(line)

with open(out_path, "w", encoding="utf-8") as f:
    f.writelines(out_lines)
PY

freeze_output=$(docker run --rm \
    --entrypoint /bin/bash \
    -v "$PWD":/work \
    -w /work \
    "$image" \
    -lc "python -m venv /tmp/venv && . /tmp/venv/bin/activate && pip install --upgrade pip setuptools >/dev/null && pip install -r /work/$(basename "$tmp_req_path") >/dev/null && pip freeze")

printf "%s\n" "$freeze_output" > "$freeze_path"

python3 - "$req_file" "$freeze_path" <<'PY'
import re
import sys

req_path = sys.argv[1]
freeze_path = sys.argv[2]

with open(req_path, "r", encoding="utf-8") as f:
    lines = f.readlines()

# Collect top-level dependency names
in_top = False
top_names = set()
for line in lines:
    stripped = line.strip()
    lower = stripped.lower()
    if lower.startswith("#") and "top level dependencies" in lower:
        in_top = True
        continue
    if lower.startswith("#") and in_top:
        in_top = False
        continue
    if not in_top:
        continue
    if not stripped or stripped.startswith("#") or stripped.startswith("-"):
        continue
    name_part = stripped.split(";", 1)[0]
    name_part = name_part.split("[", 1)[0]
    name_part = re.split(r"[<>=!~]", name_part, 1)[0]
    name_part = name_part.strip().lower().replace("_", "-")
    if name_part:
        top_names.add(name_part)

# Parse frozen dependencies
frozen = []
with open(freeze_path, "r", encoding="utf-8") as f:
    for line in f:
        stripped = line.strip()
        if not stripped or "==" not in stripped:
            continue
        name = stripped.split("==", 1)[0].strip().lower().replace("_", "-")
        if name in top_names or name == "setuptools":
            continue
        frozen.append(stripped)

frozen = sorted(set(frozen), key=str.lower)

# Replace sub dependencies section
out_lines = []
inside_sub = False
sub_header_seen = False
for i, line in enumerate(lines):
    stripped = line.strip()
    lower = stripped.lower()
    if lower.startswith("#") and ("sub depend" in lower or "sub-depend" in lower) and not sub_header_seen:
        sub_header_seen = True
        inside_sub = True
        out_lines.append(line)
        out_lines.extend([dep + "\n" for dep in frozen])
        continue
    if inside_sub:
        if lower.startswith("#") and ("sub depend" in lower or "sub-depend" in lower):
            continue
        if lower.startswith("#"):
            inside_sub = False
            if out_lines and out_lines[-1].strip():
                out_lines.append("\n")
            out_lines.append(line)
        # Skip original sub dependency lines
        continue
    out_lines.append(line)

if not sub_header_seen:
    raise SystemExit("Missing '# sub dependencies' section")

with open(req_path, "w", encoding="utf-8") as f:
    f.writelines(out_lines)
PY

echo "Updated $req_file using base image $image"