#!/usr/bin/env python3
"""Idempotent Magisk preroot hook into releasetools common.py (LOS 23.2 / A16).

Gate (either is enough):
  - ``out/.magisk/READY`` present (preferred — survives ninja/soong env scrubbing)
  - ``SUNFISH_A16_MAGISK_PREROOT=1``

Matches sfX extendrom inject (always-on once staged), adapted for A16 + our
vendor path. ``RunAndCheckOutput`` env is merged with ``os.environ`` so PATH /
ANDROID_BUILD_TOP survive (extendrom only passed ``system_root_image``; that
works via shebang but is fragile under sandboxes).
"""
from __future__ import annotations

import argparse
import os
import pathlib
import re
import sys

MARKER = "SUNFISH_A16: Magisk preroot hook"
MARKER_VER = "SUNFISH_A16: Magisk preroot hook v2"

BUILD_HOOK = '''
  # SUNFISH_A16: Magisk preroot hook v2 (before AVB)
  # Gate: staged out/.magisk/READY and/or SUNFISH_A16_MAGISK_PREROOT=1
  _sf_top = os.getenv("ANDROID_BUILD_TOP") or os.environ.get("PWD") or os.getcwd()
  _sf_ready = os.path.isfile(os.path.join(_sf_top, "out", ".magisk", "READY"))
  _sf_env = os.environ.get("SUNFISH_A16_MAGISK_PREROOT") == "1"
  if _sf_ready or _sf_env:
    if ((partition_name == "boot" and info_dict.get("init_boot") != "true")
        or partition_name == "init_boot"):
      mpatcher = [
          os.path.join(_sf_top, "vendor/sunfish-a16/magisk/patch.sh"),
          img.name,
          _sf_top,
      ]
      logger.info("SUNFISH_A16 Magisk patching (build): %s", img.name)
      _sf_penv = os.environ.copy()
      _sf_penv["system_root_image"] = info_dict.get("system_root_image", "false")
      _sf_penv.setdefault("ANDROID_BUILD_TOP", _sf_top)
      RunAndCheckOutput(mpatcher, verbose=True, env=_sf_penv)
'''

PREBUILT_HOOK = '''
  # SUNFISH_A16: Magisk preroot hook v2 (prebuilt)
  _sf_top = os.getenv("ANDROID_BUILD_TOP") or os.environ.get("PWD") or os.getcwd()
  _sf_ready = os.path.isfile(os.path.join(_sf_top, "out", ".magisk", "READY"))
  _sf_env = os.environ.get("SUNFISH_A16_MAGISK_PREROOT") == "1"
  if _sf_ready or _sf_env:
    if ((prebuilt_name == "boot.img" or prebuilt_name == "init_boot.img")
        and os.path.exists(prebuilt_path)):
      if info_dict is None:
        info_dict = OPTIONS.info_dict
      bimg = os.path.join(_sf_top, prebuilt_path)
      logger.info("SUNFISH_A16 Magisk patching (prebuilt): %s", bimg)
      mpatcher = [
          os.path.join(_sf_top, "vendor/sunfish-a16/magisk/patch.sh"),
          bimg,
          _sf_top,
      ]
      _sf_penv = os.environ.copy()
      _sf_penv["system_root_image"] = info_dict.get("system_root_image", "false")
      _sf_penv.setdefault("ANDROID_BUILD_TOP", _sf_top)
      RunAndCheckOutput(mpatcher, verbose=True, env=_sf_penv)
'''


def _strip_old_hooks(text: str) -> str:
    """Remove prior SUNFISH Magisk hook blocks (v1 gated or v2)."""
    # Build hook: from our comment through just before AVB block
    text = re.sub(
        r"\n  # SUNFISH_A16: Magisk preroot hook(?: v2)? \(before AVB\).*?"
        r"(?=\n  # AVB: if enabled)",
        "\n",
        text,
        count=1,
        flags=re.S,
    )
    # Prebuilt hook: after BOOTABLE_IMAGES assign, before "if os.path.exists(prebuilt_path)"
    text = re.sub(
        r"(  prebuilt_path = os\.path\.join\(unpack_dir, \"BOOTABLE_IMAGES\", prebuilt_name\)\n)"
        r"\n  # SUNFISH_A16: Magisk preroot hook(?: v2)? \(prebuilt\).*?"
        r"(?=\n  if os\.path\.exists\(prebuilt_path\):)",
        r"\1",
        text,
        count=1,
        flags=re.S,
    )
    return text


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--top", default=os.environ.get("ANDROID_BUILD_TOP", ""))
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument(
        "--force",
        action="store_true",
        help="Replace an existing SUNFISH Magisk hook with v2",
    )
    args = ap.parse_args()
    top = pathlib.Path(args.top or ".")
    target = top / "build/make/tools/releasetools/common.py"
    if not target.is_file():
        print(f"ERROR: missing {target}", file=sys.stderr)
        return 2
    text = target.read_text()

    if MARKER_VER in text and not args.force:
        print(f"already applied (v2): {target}")
        return 0

    if MARKER in text or MARKER_VER in text:
        if not args.force and MARKER_VER not in text:
            # old v1 present — upgrade
            pass
        text = _strip_old_hooks(text)

    needle_avb = (
        '  # AVB: if enabled, calculate and add hash to boot.img or recovery.img.\n'
        '  if info_dict.get("avb_enable") == "true":\n'
        '    avbtool = info_dict["avb_avbtool"]\n'
        '    if partition_name == "recovery":'
    )
    idx = text.find(needle_avb)
    if idx < 0:
        print("ERROR: AVB anchor not found", file=sys.stderr)
        return 3
    text = text[:idx] + BUILD_HOOK + "\n" + text[idx:]

    assign = '  prebuilt_path = os.path.join(unpack_dir, "BOOTABLE_IMAGES", prebuilt_name)\n'
    idx2 = text.find(assign)
    if idx2 < 0:
        print("ERROR: prebuilt assign not found", file=sys.stderr)
        return 4
    insert_at = idx2 + len(assign)
    text = text[:insert_at] + PREBUILT_HOOK + "\n" + text[insert_at:]

    if args.dry_run:
        print("dry-run OK (v2)")
        return 0
    target.write_text(text)
    print(f"applied v2 -> {target}")
    return 0


if __name__ == "__main__":
    sys.exit(main())