#!/usr/bin/env python3
import os
import re
import shutil
import subprocess
from pathlib import Path

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

def run(cmd, cwd: Path) -> str:
    """Run command in cwd; return stdout or die with stdout+stderr."""
    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 patch_sanity_check(patch_text: str) -> None:
    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 auto_detect_worktree(armbian_build_dir: Path) -> Path:
    root = armbian_build_dir / "cache/sources/linux-kernel-worktree"
    if not root.is_dir():
        die(f"Worktree root not found: {root}")

    # pick a directory that matches rk35xx + arm64, prefer newest by mtime
    candidates = [p for p in root.iterdir() if p.is_dir() and "rk35xx" in p.name and "arm64" in p.name]
    if not candidates:
        die(
            "Cannot auto-detect WORKTREE_DIR.\n"
            f"Searched in: {root}\n"
            "Set env WORKTREE_DIR=/full/path/to/linux-kernel-worktree/...."
        )
    candidates.sort(key=lambda p: p.stat().st_mtime, reverse=True)
    return candidates[0]

def ensure_git_repo(worktree: Path) -> None:
    out = run(["git", "rev-parse", "--is-inside-work-tree"], cwd=worktree).strip()
    if out != "true":
        die(f"Not a git worktree: {worktree}")

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 main():
    script_dir = Path(__file__).resolve().parent
    overlay_src_dir = Path(os.environ.get("OVERLAY_SRC_DIR", str(script_dir / "overlays")))

    armbian_build_dir = Path(os.environ.get("ARMBIAN_BUILD_DIR", "/home/dmn/d100/armbian/build")).expanduser()
    patch_dir = Path(os.environ.get(
        "PATCH_DIR",
        str(armbian_build_dir / "userpatches/kernel/rk35xx-vendor-6.1")
    )).expanduser()
    patch_makefile_name = os.environ.get("PATCH_MAKEFILE_NAME", "0190-rockchip-overlay-makefile.patch")
    patch_num_start = int(os.environ.get("PATCH_NUM_START", "200"))

    if not armbian_build_dir.is_dir():
        die(f"ARMBIAN_BUILD_DIR not found: {armbian_build_dir}")
    if not overlay_src_dir.is_dir():
        die(f"OVERLAY_SRC_DIR not found: {overlay_src_dir}")

    worktree_env = os.environ.get("WORKTREE_DIR", "").strip()
    if worktree_env:
        worktree = Path(worktree_env).expanduser()
    else:
        worktree = auto_detect_worktree(armbian_build_dir)

    if not worktree.is_dir():
        die(f"WORKTREE_DIR not found: {worktree}")

    ensure_git_repo(worktree)

    patch_dir.mkdir(parents=True, exist_ok=True)

    kernel_overlay_dir_rel = Path("arch/arm64/boot/dts/rockchip/overlay")
    kernel_overlay_dir = worktree / kernel_overlay_dir_rel
    kernel_overlay_dir.mkdir(parents=True, exist_ok=True)

    makefile_rel = kernel_overlay_dir_rel / "Makefile"
    makefile_path = worktree / makefile_rel
    if not makefile_path.is_file():
        die(f"Overlay Makefile not found: {makefile_path}")

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

    print(f"Worktree:  {worktree}")
    print(f"Overlays:  {overlay_src_dir}")
    print(f"Patch dir: {patch_dir}")
    print("")

    # STEP1: copy overlays into worktree
    print("STEP1: Copy overlays into kernel worktree")
    dtbos = []
    copied_relpaths = []

    for src in overlays:
        base = src.name
        name = src.stem
        dtbo = name + ".dtbo"
        dtbos.append(dtbo)

        dst_rel = kernel_overlay_dir_rel / base
        dst_abs = worktree / dst_rel

        print(f"  -> {base}")
        shutil.copy2(src, dst_abs)
        copied_relpaths.append(str(dst_rel))

        # Safe fixes for common overlay mistakes (only if present)
        txt = read_text(dst_abs)
        txt2 = re.sub(r"target-path\s*=\s*<\&", "target = <&", txt)
        txt2 = re.sub(r"^(\s*)overlay\s*\{", r"\1__overlay__ {", txt2, flags=re.M)
        if txt2 != txt:
            write_text(dst_abs, txt2)

    print("")

    # STEP2: update Makefile safely
    print("STEP2: Update overlay Makefile (sanitize + insert dtbo-y before targets)")

    mf = read_text(makefile_path)

    # sanitize glue: ".dtbo" stuck to "targets +="
    mf = re.sub(r"(\.dtbo)(targets\s+\+?=)", r"\1\n\2", mf)

    # remove wrong legacy one-liners: dtbo-$(CONFIG_ARCH_ROCKCHIP) += xxx.dtbo
    mf = re.sub(
        r"^dtbo-\$\(\s*CONFIG_ARCH_ROCKCHIP\s*\)\s*\+=\s+[^\\\n]*\.dtbo\s*$\n?",
        "",
        mf,
        flags=re.M
    )

    # find missing by EXACT line match
    missing_lines = []
    for dtbo in dtbos:
        line = f"dtbo-y += {dtbo}"
        if not re.search(rf"^\s*{re.escape(line)}\s*$", mf, flags=re.M):
            missing_lines.append(line)

    if not missing_lines:
        print("  nothing to add (sanitize applied if needed)")
    else:
        print("  adding:")
        for l in missing_lines:
            print(f"    - {l}")

        block = "".join(l + "\n" for l in missing_lines)  # important: newline after each line

        # insert before: targets += $(dtbo-y) $(dtbotxt-y)
        m = re.search(r"^targets\s+\+?=\s+\$\(\s*dtbo-y\s*\)", mf, flags=re.M)
        if m:
            mf = mf[:m.start()] + block + mf[m.start():]
        else:
            mf += "\n" + block

    # remove exact duplicate dtbo-y lines
    out = []
    seen = set()
    for line in mf.splitlines(True):
        if re.match(r"^\s*dtbo-y\s*\+=\s*[^ \t\r\n]+\.dtbo\s*$", line.strip()):
            if line in seen:
                continue
            seen.add(line)
        out.append(line)
    mf = "".join(out)

    write_text(makefile_path, mf)
    print("")

    # STEP3: Makefile patch
    print(f"STEP3: Generate Makefile patch: {patch_makefile_name}")
    run(["git", "reset", "-q"], cwd=worktree)
    run(["git", "add", str(makefile_rel)], cwd=worktree)
    changed = run(["git", "diff", "--cached", "--name-only"], cwd=worktree).splitlines()

    if str(makefile_rel) in changed:
        patch_text = run(["git", "diff", "--cached"], cwd=worktree)
        patch_sanity_check(patch_text)
        out_path = patch_dir / patch_makefile_name
        write_text(out_path, patch_text)
        print(f"  -> {out_path}")
    else:
        print("  Makefile unchanged; patch not written")
    print("")

    # STEP4: one patch per overlay .dts
    print("STEP4: Generate one patch per overlay DTS")
    n = patch_num_start
    for rel in copied_relpaths:
        base = Path(rel).name
        name = Path(base).stem
        patch_name = f"{n:04d}-rockchip-overlay-{name}.patch"

        run(["git", "reset", "-q"], cwd=worktree)
        run(["git", "add", rel], cwd=worktree)
        patch_text = run(["git", "diff", "--cached"], cwd=worktree)
        patch_sanity_check(patch_text)

        out_path = patch_dir / patch_name
        write_text(out_path, patch_text)
        print(f"  -> {out_path}")
        n += 1

    print("\nDone.")

if __name__ == "__main__":
    main()