#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
make_patches.py

Один скрипт для генерации Armbian patch-файлов из локальных DTS.

Режимы:
  --main       : генерировать патчи только для основной DTS (плата)
  --overlays   : генерировать патчи только для overlay *.dts
  --both       : генерировать и main, и overlays

ВАЖНО:
  Вместо полного пути worktree задаётся ИДЕНТИФИКАТОР:
    --worktree-id 6.1__rk35xx__arm64

  Реальный путь собирается так:
    <build-dir>/cache/sources/linux-kernel-worktree/<worktree-id>

Что делает MAIN:
  - копирует main DTS в kernel worktree:
      arch/arm64/boot/dts/<arch>/<name>.dts
  - генерирует патч DTS:
      0100-rockchip-<name>-dts.patch
  - если указан --update-makefile:
      добавляет в Makefile строку:
        dtb-$(CONFIG_ARCH_ROCKCHIP) += <name>.dtb
      (только если её там нет) и генерирует патч Makefile:
        0001-rockchip-<name>-makefile.patch

Что делает OVERLAYS:
  - копирует каждый overlay *.dts в:
      arch/arm64/boot/dts/<arch>/overlay/<name>.dts
  - добавляет в overlay/Makefile строки:
      dtbo-y += <name>.dtbo
    (без дублей, вставка перед 'targets +=')
  - генерирует патч на overlay Makefile (если были изменения):
      0190-rockchip-overlay-makefile.patch
  - генерирует отдельный патч на каждый overlay:
      0200-rockchip-overlay-<name>.patch
      0201-...

Патчи складываются сюда:
  <build-dir>/userpatches/kernel/<patch-subdir>/

Опция --clean:
  удаляет только "наши" патчи по шаблонам, чтобы не оставалось старых overlay-патчей.
"""

import argparse
import re
import shutil
import subprocess
from pathlib import Path


# -------------------------
# Вспомогательные функции
# -------------------------

def die(msg: str) -> None:
    raise SystemExit(f"ERROR: {msg}")


def run(cmd: list[str], cwd: Path) -> str:
    """
    Запускает команду и возвращает stdout.
    Если команда завершилась с ошибкой — показываем команду, cwd и вывод.
    """
    p = subprocess.run(
        cmd,
        cwd=str(cwd),
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )
    if p.returncode != 0:
        die("Command failed:\n  " + " ".join(cmd) + f"\n  (cwd={cwd})\n\nOutput:\n" + p.stdout)
    return p.stdout


def ensure_git_repo(worktree: Path) -> None:
    """Проверяем, что worktree — git репозиторий."""
    out = run(["git", "rev-parse", "--is-inside-work-tree"], cwd=worktree).strip()
    if out != "true":
        die(f"Not a git worktree: {worktree}")


def ensure_writable(path: Path) -> None:
    """
    Проверяем, что можно писать в worktree.
    Если worktree создан root'ом (docker), часто прав нет.
    """
    test = path / ".write_test.tmp"
    try:
        with open(test, "w", encoding="utf-8") as f:
            f.write("test\n")
        test.unlink(missing_ok=True)
    except PermissionError:
        die(
            f"No write permission in worktree: {path}\n\n"
            "Fix options:\n"
            "  1) Run script with sudo:\n"
            "       sudo python3 make_patches.py ...\n\n"
            "  2) Or change ownership once (recommended):\n"
            f"       sudo chown -R $USER:$USER {path}\n"
        )


def read_text(path: Path) -> str:
    return path.read_text(encoding="utf-8", errors="replace")


def write_text(path: Path, text: str) -> None:
    path.write_text(text, encoding="utf-8")


def patch_sanity_check(patch_text: str) -> None:
    """
    Контроль качества патча:
    - патч от git diff должен начинаться с 'diff --git'
    - не должен содержать мусорные '~'
    """
    if not patch_text.startswith("diff --git "):
        die("Generated patch does not start with 'diff --git' (malformed?)")
    if "\n~" in patch_text or patch_text.startswith("~"):
        die("Generated patch contains vim '~' garbage. Use only git diff output.")


def ensure_trailing_newline(text: str) -> str:
    """Makefile и патчи любят newline в конце файла."""
    return text if text.endswith("\n") else (text + "\n")


def token_in_makefile(mk_text: str, token: str) -> bool:
    """
    Проверяем, что token уже встречается в Makefile (как отдельный токен),
    чтобы не добавить дубли.
    """
    pattern = rf"(^|[\s]){re.escape(token)}($|[\s]|\\)"
    return re.search(pattern, mk_text, flags=re.M) is not None


def insert_line_before_prefix(mk_text: str, line_to_insert: str, prefix: str) -> str:
    """
    Вставляем строку line_to_insert перед первой строкой, начинающейся с prefix.
    Если prefix не нашли — добавим в конец.
    """
    mk_text = ensure_trailing_newline(mk_text)
    lines = mk_text.splitlines(keepends=True)

    out = []
    inserted = False
    for line in lines:
        if (not inserted) and line.startswith(prefix):
            out.append(line_to_insert + "\n")
            inserted = True
        out.append(line)

    if not inserted:
        out.append("\n" + line_to_insert + "\n")

    return "".join(out)


def overlay_makefile_add_dtbo_y(mk_text: str, dtbos_to_add: list[str]) -> str:
    """
    Добавляем строки:
      dtbo-y += <name>.dtbo

    Вставляем ПЕРЕД строкой:
      targets += $(dtbo-y) $(dtbotxt-y)

    Это важно для dtbs_install.
    """
    mk_text = ensure_trailing_newline(mk_text)
    lines = mk_text.splitlines(keepends=True)

    out = []
    inserted = False
    for line in lines:
        if (not inserted) and line.lstrip().startswith("targets +="):
            for d in dtbos_to_add:
                out.append(f"dtbo-y += {d}\n")
            inserted = True
        out.append(line)

    if (not inserted) and dtbos_to_add:
        out.append("\n")
        for d in dtbos_to_add:
            out.append(f"dtbo-y += {d}\n")

    return "".join(out)


def build_patch_dir(build_dir: Path, patch_subdir: str) -> Path:
    """<build>/userpatches/kernel/<subdir>"""
    return (build_dir / "userpatches/kernel" / patch_subdir).resolve()


def clean_old_patches(patch_dir: Path, do_main: bool, do_overlays: bool) -> None:
    """
    Удаляем только наши патчи (по шаблонам), чтобы не ломать другие патчи.

    MAIN:
      0001-rockchip-*-makefile.patch
      0100-rockchip-*-dts.patch

    OVERLAYS:
      0190-rockchip-overlay-makefile.patch
      02??-rockchip-overlay-*.patch   (0200..0299)
    """
    patterns = []
    if do_main:
        patterns += [
            "0001-rockchip-*-makefile.patch",
            "0100-rockchip-*-dts.patch",
        ]
    if do_overlays:
        patterns += [
            "0190-rockchip-overlay-makefile.patch",
            "02??-rockchip-overlay-*.patch",
        ]

    removed = 0
    for pat in patterns:
        for p in patch_dir.glob(pat):
            try:
                p.unlink()
                removed += 1
            except OSError as e:
                die(f"Cannot remove {p}: {e}")

    if removed:
        print(f"[CLEAN] Removed {removed} old patch file(s) from {patch_dir}")
    else:
        print(f"[CLEAN] No old patch files matched in {patch_dir}")


# -------------------------
# Основная программа
# -------------------------

def main():
    ap = argparse.ArgumentParser(
        description="Generate Armbian patches for main DTS and/or overlays from local sources."
    )

    # Режим работы (один из)
    mode = ap.add_mutually_exclusive_group(required=True)
    mode.add_argument("--main", action="store_true", help="Generate patches only for the main DTS")
    mode.add_argument("--overlays", action="store_true", help="Generate patches only for overlays directory")
    mode.add_argument("--both", action="store_true", help="Generate patches for both main DTS and overlays")

    # Общие параметры
    ap.add_argument("--build-dir", default="/home/dmn/d100/armbian/build", help="Armbian build dir")
    ap.add_argument("--worktree-id", required=True, help="Kernel worktree ID, e.g. 6.1__rk35xx__arm64 (REQUIRED)")
    ap.add_argument("--patch-subdir", required=True, help="Subfolder inside <build-dir>/userpatches/kernel/, e.g. rk35xx-vendor-6.1 or rk35xx-current")
    ap.add_argument("--arch", default="rockchip", help="DTS arch subdir under arch/arm64/boot/dts (default: rockchip)")
    ap.add_argument("--clean", action="store_true", help="Remove previously generated patch files (safe patterns) before writing new ones")

    # MAIN DTS параметры
    ap.add_argument("--main-dts", default="", help="Path to main DTS file (host), e.g. ./rk3568-napi2.dts")
    ap.add_argument("--main-dts-patch", default="", help="Main DTS patch filename (default: 0100-rockchip-<name>-dts.patch)")
    ap.add_argument("--main-makefile-patch", default="", help="Main Makefile patch filename (default: 0001-rockchip-<name>-makefile.patch)")
    ap.add_argument("--update-makefile", action="store_true", help="For main DTS: ensure dtb line exists in Makefile and create Makefile patch if changed")

    # OVERLAYS параметры
    ap.add_argument("--overlays-dir", default="", help="Directory containing overlay *.dts files (host), e.g. ./overlays")
    ap.add_argument("--overlay-makefile-patch", default="0190-rockchip-overlay-makefile.patch", help="Patch name for overlay/Makefile changes")
    ap.add_argument("--overlay-patch-start", type=int, default=200, help="First number for per-overlay patches (default: 200)")

    args = ap.parse_args()

    # build-dir
    build_dir = Path(args.build_dir).expanduser().resolve()
    if not build_dir.is_dir():
        die(f"--build-dir not found: {build_dir}")

    # worktree path from id (NO auto-detect!)
    worktree = (build_dir / "cache/sources/linux-kernel-worktree" / args.worktree_id).resolve()
    if not worktree.is_dir():
        die(
            f"Kernel worktree not found:\n"
            f"  {worktree}\n\n"
            f"Expected layout:\n"
            f"  <build-dir>/cache/sources/linux-kernel-worktree/<worktree-id>\n"
            f"Given worktree-id:\n"
            f"  {args.worktree_id}"
        )

    # validate worktree
    ensure_git_repo(worktree)
    ensure_writable(worktree)

    # patch dir
    patch_dir = build_patch_dir(build_dir, args.patch_subdir)
    patch_dir.mkdir(parents=True, exist_ok=True)

    do_main = args.main or args.both
    do_overlays = args.overlays or args.both

    # optional cleanup
    if args.clean:
        clean_old_patches(patch_dir, do_main, do_overlays)

    # summary
    print(f"Worktree:   {worktree}")
    print(f"Build dir:  {build_dir}")
    print(f"Patch dir:  {patch_dir}")
    print(f"Arch:       {args.arch}")
    print("")

    # -----------------------
    # MAIN DTS
    # -----------------------
    if do_main:
        if not args.main_dts:
            die("--main-dts is required for --main or --both")

        src = Path(args.main_dts).expanduser().resolve()
        if not src.is_file():
            die(f"--main-dts not found: {src}")

        main_dts_rel = Path("arch/arm64/boot/dts") / args.arch / src.name
        main_makefile_rel = main_dts_rel.parent / "Makefile"

        name = src.stem  # например rk3568-napi2
        dts_patch_name = args.main_dts_patch or f"0100-rockchip-{name}-dts.patch"
        mk_patch_name = args.main_makefile_patch or f"0001-rockchip-{name}-makefile.patch"

        dst_abs = worktree / main_dts_rel
        dst_abs.parent.mkdir(parents=True, exist_ok=True)

        print("=== MAIN DTS ===")
        print(f"SRC: {src}")
        print(f"DST: {main_dts_rel}")

        print("[1/3] Copy main DTS into kernel worktree")
        shutil.copy2(src, dst_abs)

        print("[2/3] Generate main DTS patch")
        run(["git", "reset", "-q"], cwd=worktree)
        run(["git", "add", str(main_dts_rel)], cwd=worktree)
        dts_patch = run(["git", "diff", "--cached"], cwd=worktree)
        patch_sanity_check(dts_patch)

        if str(main_dts_rel) not in dts_patch:
            die(f"Main DTS patch does not reference {main_dts_rel}")

        out_dts = patch_dir / dts_patch_name
        write_text(out_dts, dts_patch)
        print(f"      -> {out_dts}")

        print("[3/3] (optional) Update main Makefile + Makefile patch")
        if args.update_makefile:
            makefile_abs = worktree / main_makefile_rel
            if not makefile_abs.is_file():
                die(f"Main Makefile not found: {makefile_abs}")

            dtb_name = f"{name}.dtb"
            dtb_line = f"dtb-$(CONFIG_ARCH_ROCKCHIP) += {dtb_name}"

            mk = read_text(makefile_abs)

            if token_in_makefile(mk, dtb_name):
                print("      dtb already present in main Makefile")
            else:
                print("      inserting dtb line before 'subdir-y'")
                mk2 = insert_line_before_prefix(mk, dtb_line, "subdir-y")
                write_text(makefile_abs, mk2)

            run(["git", "reset", "-q"], cwd=worktree)
            run(["git", "add", str(main_makefile_rel)], cwd=worktree)
            changed = run(["git", "diff", "--cached", "--name-only"], cwd=worktree).splitlines()

            if str(main_makefile_rel) in changed:
                mk_patch = run(["git", "diff", "--cached"], cwd=worktree)
                patch_sanity_check(mk_patch)
                out_mk = patch_dir / mk_patch_name
                write_text(out_mk, mk_patch)
                print(f"      -> {out_mk}")
            else:
                print("      Makefile unchanged; no Makefile patch written")
        else:
            print("      skipped (use --update-makefile if needed)")

        print("")

    # -----------------------
    # OVERLAYS
    # -----------------------
    if do_overlays:
        if not args.overlays_dir:
            die("--overlays-dir is required for --overlays or --both")

        overlays_dir = Path(args.overlays_dir).expanduser().resolve()
        if not overlays_dir.is_dir():
            die(f"--overlays-dir not found: {overlays_dir}")

        overlays = sorted(overlays_dir.glob("*.dts"))
        if not overlays:
            die(f"No *.dts overlays found in: {overlays_dir}")

        kernel_overlay_dir_rel = Path("arch/arm64/boot/dts") / args.arch / "overlay"
        kernel_overlay_makefile_rel = kernel_overlay_dir_rel / "Makefile"

        kernel_overlay_dir_abs = worktree / kernel_overlay_dir_rel
        kernel_overlay_makefile_abs = worktree / kernel_overlay_makefile_rel

        print("=== OVERLAYS ===")
        print(f"SRC DIR: {overlays_dir}")
        print(f"DST DIR: {kernel_overlay_dir_rel}")

        print("[1/3] Copy overlay DTS files into kernel worktree")
        kernel_overlay_dir_abs.mkdir(parents=True, exist_ok=True)

        dtbos = []
        for src in overlays:
            dst_rel = kernel_overlay_dir_rel / src.name
            dst_abs = worktree / dst_rel
            shutil.copy2(src, dst_abs)
            dtbos.append(src.stem + ".dtbo")
            print(f"      copied {src.name}")

        print("[2/3] Update overlay Makefile once (dedup) + Makefile patch if changed")
        if not kernel_overlay_makefile_abs.is_file():
            die(f"Overlay Makefile not found: {kernel_overlay_makefile_abs}")

        mk = ensure_trailing_newline(read_text(kernel_overlay_makefile_abs))
        to_add = [d for d in dtbos if not token_in_makefile(mk, d)]

        if not to_add:
            print("      nothing to add")
        else:
            mk2 = overlay_makefile_add_dtbo_y(mk, to_add)
            write_text(kernel_overlay_makefile_abs, mk2)

            print("      added:")
            for d in to_add:
                print(f"        dtbo-y += {d}")

            run(["git", "reset", "-q"], cwd=worktree)
            run(["git", "add", str(kernel_overlay_makefile_rel)], cwd=worktree)
            mk_patch = run(["git", "diff", "--cached"], cwd=worktree)
            patch_sanity_check(mk_patch)

            out_mk_patch = patch_dir / args.overlay_makefile_patch
            write_text(out_mk_patch, mk_patch)
            print(f"      -> {out_mk_patch}")

        print("[3/3] Generate one patch per overlay DTS")
        n = args.overlay_patch_start

        for src in overlays:
            name = src.stem
            dst_rel = kernel_overlay_dir_rel / src.name

            run(["git", "reset", "-q"], cwd=worktree)
            run(["git", "add", str(dst_rel)], cwd=worktree)
            patch = run(["git", "diff", "--cached"], cwd=worktree)
            patch_sanity_check(patch)

            if str(dst_rel) not in patch:
                die(f"Overlay patch does not reference {dst_rel}")

            patch_name = f"{n:04d}-rockchip-overlay-{name}.patch"
            out_patch = patch_dir / patch_name
            write_text(out_patch, patch)
            print(f"      -> {out_patch}")
            n += 1

        print("")

    print("Done.")


if __name__ == "__main__":
    main()