#!/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 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 auto_detect_worktree(build_dir: Path) -> Path:
    root = build_dir / "cache/sources/linux-kernel-worktree"
    if not root.is_dir():
        die(f"Worktree root not found: {root}")

    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/...."
        )
    # prefer newest by mtime
    candidates.sort(key=lambda p: p.stat().st_mtime, reverse=True)
    return candidates[0]

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:
    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 makefile_has_dtb_token(mk_text: str, dtb_name: str) -> bool:
    # Token match, like: (^|space)rk3568-napi2.dtb($|space|\)
    pattern = rf"(^|[\s]){re.escape(dtb_name)}($|[\s]|\\)"
    return re.search(pattern, mk_text, flags=re.M) is not None

def ensure_writable(path: Path) -> None:
    """Check we can write inside path; if not, print actionable fix."""
    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 with sudo:\n"
            "       sudo python3 make-main-dts.py\n\n"
            "  2) Or make worktree owned by your user (recommended):\n"
            f"       sudo chown -R $USER:$USER {path}\n"
        )

def insert_dtb_line_before_subdir_y(mk_text: str, dtb_line: str) -> str:
    """
    Insert dtb line before the first 'subdir-y' line.
    If not found, append at end (fallback).
    """
    lines = mk_text.splitlines(keepends=True)
    out = []
    inserted = False
    for line in lines:
        if (not inserted) and line.startswith("subdir-y"):
            out.append(dtb_line + "\n")
            inserted = True
        out.append(line)
    if not inserted:
        # fallback: append once
        if not mk_text.endswith("\n"):
            out.append("\n")
        out.append(dtb_line + "\n")
    return "".join(out)

def main():
    script_dir = Path(__file__).resolve().parent

    # env / defaults
    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()

    src_dts = Path(os.environ.get("SRC_DTS", str(script_dir / "rk3568-napi2.dts"))).expanduser()

    kernel_dts_rel = Path("arch/arm64/boot/dts/rockchip/rk3568-napi2.dts")
    kernel_makefile_rel = Path("arch/arm64/boot/dts/rockchip/Makefile")

    patch_makefile_name = os.environ.get("PATCH_MAKEFILE_NAME", "0001-rockchip-rk3568-napi2-makefile.patch")
    patch_dts_name = os.environ.get("PATCH_DTS_NAME", "0100-rockchip-rk3568-napi2-dts.patch")

    if not armbian_build_dir.is_dir():
        die(f"ARMBIAN_BUILD_DIR not found: {armbian_build_dir}")
    if not src_dts.is_file():
        die(f"SRC_DTS not found: {src_dts}")

    worktree_env = os.environ.get("WORKTREE_DIR", "").strip()
    worktree = Path(worktree_env).expanduser() if worktree_env else auto_detect_worktree(armbian_build_dir)
    if not worktree.is_dir():
        die(f"WORKTREE_DIR not found: {worktree}")

    ensure_git_repo(worktree)
    ensure_writable(worktree)

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

    kernel_dts_abs = worktree / kernel_dts_rel
    kernel_makefile_abs = worktree / kernel_makefile_rel
    if not kernel_makefile_abs.is_file():
        die(f"Kernel Makefile not found: {kernel_makefile_abs}")

    print(f"Worktree:  {worktree}")
    print(f"Patch dir: {patch_dir}")
    print(f"SRC DTS:   {src_dts}\n")

    # [1/4] Copy DTS
    print("[1/4] Copy main DTS into kernel worktree")
    kernel_dts_abs.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(src_dts, kernel_dts_abs)

    # [2/4] Ensure Makefile contains dtb line (insert at correct position)
    print("[2/4] Ensure main Makefile contains rk3568-napi2.dtb (only if missing)")
    mk = read_text(kernel_makefile_abs)
    dtb_name = "rk3568-napi2.dtb"
    dtb_line = f"dtb-$(CONFIG_ARCH_ROCKCHIP) += {dtb_name}"

    if makefile_has_dtb_token(mk, dtb_name):
        print("      already present")
    else:
        print("      inserting dtb line before subdir-y")
        mk2 = insert_dtb_line_before_subdir_y(mk, dtb_line)
        write_text(kernel_makefile_abs, mk2)

    # [3/4] Generate DTS patch
    print("[3/4] Generate DTS patch")
    run(["git", "reset", "-q"], cwd=worktree)
    run(["git", "add", str(kernel_dts_rel)], cwd=worktree)
    dts_patch = run(["git", "diff", "--cached"], cwd=worktree)
    patch_sanity_check(dts_patch)
    if str(kernel_dts_rel) not in dts_patch:
        die(f"DTS patch does not reference {kernel_dts_rel}")
    out_dts = patch_dir / patch_dts_name
    write_text(out_dts, dts_patch)
    print(f"      -> {out_dts}")

    # [4/4] Generate Makefile patch only if Makefile changed
    print("[4/4] Generate Makefile patch (only if Makefile changed)")
    run(["git", "reset", "-q"], cwd=worktree)
    run(["git", "add", str(kernel_makefile_rel)], cwd=worktree)
    changed = run(["git", "diff", "--cached", "--name-only"], cwd=worktree).splitlines()

    if str(kernel_makefile_rel) in changed:
        mk_patch = run(["git", "diff", "--cached"], cwd=worktree)
        patch_sanity_check(mk_patch)

        # extra safety: if patch accidentally contains a lot of duplicated Armbian comments, warn
        if mk_patch.count("Armbian: Incremental: assuming overlay targets are already in the Makefile") > 2:
            die("Makefile patch contains duplicated Armbian comments. Makefile is likely polluted; reset it and rerun.")

        out_mk = patch_dir / patch_makefile_name
        write_text(out_mk, mk_patch)
        print(f"      saved {out_mk.name}")
        print(f"      -> {out_mk}")
    else:
        print("      Makefile unchanged, skip")

    print(f"\nDone. Patches in: {patch_dir}")

if __name__ == "__main__":
    main()