#!/usr/bin/env python3
import argparse
import hashlib
import json
import os
import re
import shutil
import stat as pystat
import subprocess
import sys
import time
from pathlib import Path, PurePosixPath

LV = "/dev/mapper/plex--new-plex--new"
DEST_MOUNT = Path("/mnt/recovery")
DEAD_START = 3418636288
DEAD_END = 4395389951
RESERVE_BYTES = 2 * 1024**3

extent_re = re.compile(r"\((\d+)(?:-(\d+))?\):(\d+)(?:-(\d+))?")
etb_re = re.compile(r"\(ETB(\d+)\):(\d+)")
inode_re = re.compile(r"Inode:\s*(\d+)\s+Type:\s*([^\s]+)")
size_re = re.compile(r"\bSize:\s*(\d+)")
blockcount_re = re.compile(r"\bBlockcount:\s*(\d+)")

def die(msg):
    print(f"ERROR: {msg}", file=sys.stderr)
    raise SystemExit(1)

def human(n):
    units = ["B", "KiB", "MiB", "GiB", "TiB"]
    x = float(n)
    for unit in units:
        if x < 1024 or unit == units[-1]:
            return f"{x:.2f} {unit}"
        x /= 1024

def run(cmd):
    return subprocess.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

def debugfs(command):
    env = os.environ.copy()
    env["DEBUGFS_PAGER"] = "cat"
    return subprocess.run(
        ["debugfs", "-c", "-R", command, LV],
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env=env,
    )

def normalize_source(source):
    source = source.strip()
    if not source:
        die("Source path cannot be empty.")
    if not source.startswith("/"):
        source = "/" + source

    parts = []
    for part in PurePosixPath(source).parts:
        if part in ("/", "", "."):
            continue
        if part == "..":
            die("Source path may not contain '..'.")
        parts.append(part)

    if not parts:
        die("Refusing to scan filesystem root '/'. Specify a subtree.")

    return "/" + "/".join(parts)

def quote_debugfs_path(path):
    return '"' + path.replace("\\", "\\\\").replace('"', '\\"') + '"'

def slug_for(source):
    base = source.strip("/").lower()
    base = re.sub(r"[^a-z0-9]+", "-", base).strip("-") or "root"
    base = base[:80].rstrip("-")
    digest = hashlib.sha1(source.encode("utf-8")).hexdigest()[:8]
    return f"{base}-{digest}"

def config_for(source):
    source = normalize_source(source)
    rel = PurePosixPath(source.lstrip("/"))
    slug = slug_for(source)
    return {
        "source": source,
        "dest": DEST_MOUNT.joinpath(*rel.parts),
        "slug": slug,
        "manifest": DEST_MOUNT / f"salvage-{slug}-manifest.jsonl",
        "scan_summary": DEST_MOUNT / f"salvage-{slug}-scan-summary.txt",
        "recovery_log": DEST_MOUNT / f"salvage-{slug}-recovery.tsv",
        "recovery_summary": DEST_MOUNT / f"salvage-{slug}-recovery-summary.txt",
        "tmp_dir": DEST_MOUNT / f".salvage-{slug}-tmp",
    }

def preflight(cfg, require_manifest=False):
    if os.geteuid() != 0:
        die("Run as root.")

    for cmd in ("debugfs", "blockdev", "findmnt"):
        if shutil.which(cmd) is None:
            die(f"Required command not found: {cmd}")

    if not Path(LV).exists():
        die(f"LV does not exist: {LV}")

    ro = run(["blockdev", "--getro", LV])
    if ro.returncode != 0 or ro.stdout.strip() != "1":
        die(f"{LV} is not kernel read-only. Refusing to continue.")

    mount = run(["findmnt", "-n", "-o", "TARGET,SOURCE,FSTYPE", "-T", str(DEST_MOUNT)])
    if mount.returncode != 0:
        die(f"{DEST_MOUNT} is not mounted.")

    fields = mount.stdout.strip().split(None, 2)
    if not fields or fields[0] != str(DEST_MOUNT):
        die(f"{DEST_MOUNT} is not the mount point itself.")

    if len(fields) >= 2 and "plex--new-plex--new" in fields[1]:
        die("Recovery destination is on the damaged LV.")

    if require_manifest and not cfg["manifest"].exists():
        die(f"Manifest not found: {cfg['manifest']}. Run scan first.")

    print("Pre-flight:")
    print(f"  Source LV:       {LV} (read-only)")
    print(f"  Source subtree:  {cfg['source']}")
    print(f"  Destination:     {cfg['dest']}")
    print(f"  Dead blocks:     {DEAD_START}-{DEAD_END}")
    print(f"  Recovery mount:  {mount.stdout.strip()}")
    print(f"  Job slug:        {cfg['slug']}")
    print()

def source_root_inode(source):
    result = debugfs(f"stat {quote_debugfs_path(source)}")
    match = inode_re.search(result.stdout)
    if not match or match.group(2) != "directory":
        detail = (result.stdout + "\n" + result.stderr).strip()
        die(f"Could not open {source} as a directory.\n{detail}")
    return int(match.group(1))

def list_directory(inode):
    result = debugfs(f"ls -p <{inode}>")
    entries = []
    for raw in result.stdout.splitlines():
        if not raw.startswith("/"):
            continue
        parts = raw.split("/")
        if len(parts) < 7:
            continue
        try:
            child_inode = int(parts[1])
            mode = int(parts[2], 8)
        except ValueError:
            continue
        name = parts[5]
        if name in (".", "..") or child_inode == 0:
            continue
        entries.append((child_inode, mode, name))
    return entries, result

def inspect_regular_file(inode):
    result = debugfs(f"stat <{inode}>")
    text = result.stdout
    im = inode_re.search(text)
    sm = size_re.search(text)
    bcm = blockcount_re.search(text)

    if not im or im.group(2) != "regular" or not sm:
        return {"status": "UNKNOWN", "reason": "stat output could not be parsed as a regular file"}

    size = int(sm.group(1))
    blockcount = int(bcm.group(1)) if bcm else None

    if size == 0:
        return {"status": "SAFE", "reason": "zero-length file", "size": 0, "blockcount": blockcount, "extents": []}

    extent_text = text.split("EXTENTS:", 1)[1] if "EXTENTS:" in text else ""
    extents = []
    for m in extent_re.finditer(extent_text):
        extents.append({
            "logical_start": int(m.group(1)),
            "logical_end": int(m.group(2) or m.group(1)),
            "physical_start": int(m.group(3)),
            "physical_end": int(m.group(4) or m.group(3)),
        })

    etbs = [{"level": int(level), "block": int(block)} for level, block in etb_re.findall(extent_text)]

    if extents:
        overlaps = [
            ex for ex in extents
            if ex["physical_start"] <= DEAD_END and ex["physical_end"] >= DEAD_START
        ]
        if overlaps:
            return {
                "status": "LOST",
                "reason": "one or more data extents overlap the dead PV",
                "size": size,
                "blockcount": blockcount,
                "extents": extents,
                "overlaps": overlaps,
                "etbs": etbs,
            }

        return {
            "status": "SAFE",
            "reason": "all physical data extents are outside the dead PV",
            "size": size,
            "blockcount": blockcount,
            "extents": extents,
            "etbs": etbs,
        }

    if etbs:
        if any(DEAD_START <= e["block"] <= DEAD_END for e in etbs):
            return {
                "status": "ETB_LOST",
                "reason": "extent-tree metadata block is on the dead PV",
                "size": size,
                "blockcount": blockcount,
                "etbs": etbs,
            }
        return {
            "status": "ETB_SURVIVES",
            "reason": "extent-tree metadata is outside the dead PV; investigate",
            "size": size,
            "blockcount": blockcount,
            "etbs": etbs,
        }

    return {
        "status": "UNKNOWN",
        "reason": "no parsable data extents or extent-tree metadata references",
        "size": size,
        "blockcount": blockcount,
    }

def scan(source):
    cfg = config_for(source)
    preflight(cfg)
    root_inode = source_root_inode(cfg["source"])

    print(f"Scanning {cfg['source']} from inode {root_inode}...")
    print("No file data will be copied during this stage.\n")

    tmp_manifest = cfg["manifest"].with_suffix(".jsonl.new")
    if tmp_manifest.exists():
        tmp_manifest.unlink()

    statuses = ["SAFE", "LOST", "ETB_LOST", "ETB_SURVIVES", "UNKNOWN", "OTHER"]
    counts = {s: 0 for s in statuses}
    sizes = {s: 0 for s in statuses if s != "OTHER"}
    dirs = 0
    files_seen = 0
    started = time.time()
    stack = [(root_inode, Path("."))]

    with tmp_manifest.open("w", encoding="utf-8") as mf:
        mf.write(json.dumps({
            "_meta": {
                "source": cfg["source"],
                "dest": str(cfg["dest"]),
                "slug": cfg["slug"],
                "dead_start": DEAD_START,
                "dead_end": DEAD_END,
            }
        }) + "\n")

        while stack:
            dir_inode, rel_dir = stack.pop()
            dirs += 1
            entries, result = list_directory(dir_inode)

            if not entries and result.stderr.strip():
                print(f"WARN directory inode {dir_inode} ({rel_dir}): {result.stderr.strip().splitlines()[-1]}")

            for inode, mode, name in entries:
                rel_path = rel_dir / name
                rel_str = str(rel_path).removeprefix("./")

                if pystat.S_ISDIR(mode):
                    stack.append((inode, rel_path))
                    continue

                if not pystat.S_ISREG(mode):
                    counts["OTHER"] += 1
                    mf.write(json.dumps({
                        "status": "OTHER",
                        "inode": inode,
                        "path": rel_str,
                        "mode": oct(mode),
                        "reason": "not a regular file",
                    }, ensure_ascii=False) + "\n")
                    mf.flush()
                    continue

                files_seen += 1
                info = inspect_regular_file(inode)
                status = info["status"]
                counts[status] += 1
                if "size" in info:
                    sizes[status] += info["size"]

                mf.write(json.dumps({
                    "status": status,
                    "inode": inode,
                    "path": rel_str,
                    **info,
                }, ensure_ascii=False) + "\n")
                mf.flush()

                if status != "SAFE":
                    print(f"{status:<12} {rel_str}")

                if files_seen % 100 == 0:
                    elapsed = time.time() - started
                    print(
                        f"[{files_seen} files] SAFE {counts['SAFE']} ({human(sizes['SAFE'])}), "
                        f"LOST {counts['LOST']} ({human(sizes['LOST'])}), "
                        f"ETB_LOST {counts['ETB_LOST']}, elapsed {elapsed/60:.1f} min"
                    )

    tmp_manifest.replace(cfg["manifest"])

    lines = [
        "GENERIC PLEX SALVAGE SCAN",
        "=========================",
        f"Source:       {cfg['source']}",
        f"Destination:  {cfg['dest']}",
        f"Root inode:   {root_inode}",
        f"Dead blocks:  {DEAD_START}-{DEAD_END}",
        f"Directories:  {dirs}",
        f"Files seen:   {files_seen}",
        "",
    ]
    for status in ["SAFE", "LOST", "ETB_LOST", "ETB_SURVIVES", "UNKNOWN"]:
        lines.append(f"{status + ':':14} {counts[status]} files, {human(sizes[status])}")
    lines += [
        f"{'OTHER:':14} {counts['OTHER']} entries",
        "",
        f"Manifest:     {cfg['manifest']}",
    ]
    summary = "\n".join(lines) + "\n"
    cfg["scan_summary"].write_text(summary, encoding="utf-8")
    print("\n" + summary)

def load_manifest(cfg):
    records = []
    meta = None
    with cfg["manifest"].open("r", encoding="utf-8") as f:
        for line_no, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                rec = json.loads(line)
            except json.JSONDecodeError as exc:
                die(f"Bad manifest line {line_no}: {exc}")
            if "_meta" in rec:
                meta = rec["_meta"]
            else:
                records.append(rec)

    if not meta:
        die("Manifest has no job metadata header.")
    if meta.get("source") != cfg["source"]:
        die(f"Manifest source mismatch: {meta.get('source')} != {cfg['source']}")
    return records

def append_log(path, status, size, inode, relpath, detail=""):
    new_file = not path.exists()
    with path.open("a", encoding="utf-8") as f:
        if new_file:
            f.write("status\tbytes\tinode\tpath\tdetail\n")
        detail = detail.replace("\t", " ").replace("\n", " ")
        f.write(f"{status}\t{size}\t{inode}\t{relpath}\t{detail}\n")
        f.flush()

def recover(source):
    cfg = config_for(source)
    preflight(cfg, require_manifest=True)
    records = load_manifest(cfg)
    safe_records = [r for r in records if r.get("status") == "SAFE"]

    safe_bytes = sum(int(r.get("size", 0)) for r in safe_records)
    usage = shutil.disk_usage(DEST_MOUNT)

    print(f"Manifest contains {len(safe_records)} SAFE files totaling {human(safe_bytes)}.")
    print(f"Recovery drive free space: {human(usage.free)}")
    print(f"Reserved free space: {human(RESERVE_BYTES)}\n")

    cfg["dest"].mkdir(parents=True, exist_ok=True)
    cfg["tmp_dir"].mkdir(parents=True, exist_ok=True)

    recovered = recovered_bytes = existing = existing_bytes = errors = 0
    started = time.time()

    for idx, rec in enumerate(safe_records, 1):
        inode = int(rec["inode"])
        size = int(rec.get("size", 0))
        rel = Path(rec["path"])
        dest = cfg["dest"] / rel

        try:
            dest.parent.mkdir(parents=True, exist_ok=True)
        except OSError as exc:
            append_log(cfg["recovery_log"], "DEST_ERROR", size, inode, str(rel), str(exc))
            errors += 1
            continue

        if dest.exists():
            try:
                if dest.is_file() and dest.stat().st_size == size:
                    existing += 1
                    existing_bytes += size
                    append_log(cfg["recovery_log"], "EXISTS", size, inode, str(rel), "existing size matches inode")
                    continue
            except OSError:
                pass

        free = shutil.disk_usage(DEST_MOUNT).free
        if free < size + RESERVE_BYTES:
            msg = f"need {human(size)} plus {human(RESERVE_BYTES)} reserve; only {human(free)} free"
            append_log(cfg["recovery_log"], "NO_SPACE", size, inode, str(rel), msg)
            print(f"Stopping before {rel}: {msg}")
            break

        tmp = cfg["tmp_dir"] / f"{inode}.part"
        try:
            if tmp.exists():
                tmp.unlink()
        except OSError as exc:
            append_log(cfg["recovery_log"], "TMP_ERROR", size, inode, str(rel), str(exc))
            errors += 1
            continue

        result = debugfs(f"dump <{inode}> {tmp}")

        if not tmp.exists():
            detail = (result.stdout + "\n" + result.stderr).strip()
            append_log(cfg["recovery_log"], "DUMP_ERROR", size, inode, str(rel), detail)
            errors += 1
            print(f"DUMP_ERROR {rel}")
            continue

        actual = tmp.stat().st_size
        if actual != size:
            append_log(cfg["recovery_log"], "SIZE_MISMATCH", size, inode, str(rel), f"dumped {actual} bytes; inode says {size}")
            errors += 1
            try:
                tmp.unlink()
            except OSError:
                pass
            print(f"SIZE_MISMATCH {rel}: {actual} != {size}")
            continue

        try:
            os.replace(tmp, dest)
        except OSError as exc:
            append_log(cfg["recovery_log"], "MOVE_ERROR", size, inode, str(rel), str(exc))
            errors += 1
            continue

        recovered += 1
        recovered_bytes += size
        append_log(cfg["recovery_log"], "RECOVERED", size, inode, str(rel))
        print(f"RECOVERED [{idx}/{len(safe_records)}] {human(size):>10}  {rel}")

        if recovered % 50 == 0:
            elapsed = time.time() - started
            print(f"  -> {recovered} recovered, {human(recovered_bytes)}, {errors} errors, {elapsed/60:.1f} min")

    free_final = shutil.disk_usage(DEST_MOUNT).free
    summary = (
        "GENERIC PLEX SALVAGE RECOVERY\n"
        "=============================\n"
        f"Source:          {cfg['source']}\n"
        f"Destination:     {cfg['dest']}\n"
        f"Recovered:       {recovered} files, {human(recovered_bytes)}\n"
        f"Already present: {existing} files, {human(existing_bytes)}\n"
        f"Errors:          {errors}\n"
        f"Free remaining:  {human(free_final)}\n"
        f"Log:             {cfg['recovery_log']}\n"
    )
    cfg["recovery_summary"].write_text(summary, encoding="utf-8")
    print("\n" + summary)

def main():
    parser = argparse.ArgumentParser(description="Generic read-only debugfs salvage tool")
    parser.add_argument("mode", choices=("scan", "recover"))
    parser.add_argument("source", help="source directory inside the damaged filesystem")
    args = parser.parse_args()

    if args.mode == "scan":
        scan(args.source)
    else:
        recover(args.source)

if __name__ == "__main__":
    main()