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

LV = "/dev/mapper/plex--new-plex--new"
MANIFEST = Path("/mnt/recovery/plex-books-manifest.jsonl")
REPORT = Path("/mnt/recovery/plex-books-unknown-check.tsv")
SUMMARY = Path("/mnt/recovery/plex-books-unknown-summary.txt")

DEAD_START = 3418636288
DEAD_END = 4395389951

inode_re = re.compile(r"Inode:\s*(\d+)\s+Type:\s*([^\s]+)")
size_re = re.compile(r"\bSize:\s*(\d+)")
etb_re = re.compile(r"\(ETB(\d+)\):(\d+)")
extent_re = re.compile(r"\((\d+)(?:-(\d+))?\):(\d+)(?:-(\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 preflight():
    if os.geteuid() != 0:
        die("Run as root.")
    if not Path(LV).exists():
        die(f"Missing LV: {LV}")
    ro = run(["blockdev", "--getro", LV])
    if ro.returncode != 0 or ro.stdout.strip() != "1":
        die(f"{LV} is not kernel read-only.")
    if not MANIFEST.exists():
        die(f"Manifest not found: {MANIFEST}")

def classify(record):
    inode = int(record["inode"])
    path = record["path"]
    size = int(record.get("size", 0))

    stat_result = debugfs(f"stat <{inode}>")
    stat_text = stat_result.stdout

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

    real_extents = []
    if "EXTENTS:" in stat_text:
        extent_text = stat_text.split("EXTENTS:", 1)[1]
        for m in extent_re.finditer(extent_text):
            # Skip ETB pseudo-lines; extent_re only matches numeric logical starts.
            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))
            real_extents.append(
                {
                    "logical_start": logical_start,
                    "logical_end": logical_end,
                    "physical_start": physical_start,
                    "physical_end": physical_end,
                }
            )

    blocks_result = debugfs(f"blocks <{inode}>")
    block_nums = []
    for token in re.findall(r"\b\d+\b", blocks_result.stdout):
        try:
            block_nums.append(int(token))
        except ValueError:
            pass

    dead_etbs = [e for e in etbs if DEAD_START <= e["block"] <= DEAD_END]
    live_etbs = [e for e in etbs if not (DEAD_START <= e["block"] <= DEAD_END)]

    overlapping_extents = [
        ex for ex in real_extents
        if ex["physical_start"] <= DEAD_END and ex["physical_end"] >= DEAD_START
    ]

    if real_extents:
        if overlapping_extents:
            cls = "DATA_OVERLAP"
            reason = "parsable data extent overlaps dead PV"
        else:
            cls = "EXTENTS_RECOVERABLE"
            reason = "parsable data extents exist and avoid dead PV"
    elif etbs:
        if dead_etbs:
            cls = "ETB_LOST"
            reason = "extent-tree metadata block is on dead PV"
        elif live_etbs:
            cls = "ETB_SURVIVES"
            reason = "extent-tree metadata block is outside dead PV; investigate"
        else:
            cls = "OTHER_UNKNOWN"
            reason = "ETB metadata present but could not classify"
    else:
        cls = "OTHER_UNKNOWN"
        reason = "no parsable data extents or ETB references"

    return {
        "class": cls,
        "reason": reason,
        "inode": inode,
        "path": path,
        "size": size,
        "etbs": etbs,
        "dead_etbs": dead_etbs,
        "live_etbs": live_etbs,
        "real_extents": real_extents,
        "blocks_output": " ".join(str(b) for b in block_nums),
        "stat_stderr": stat_result.stderr.strip(),
        "blocks_stderr": blocks_result.stderr.strip(),
    }

def main():
    preflight()

    unknown = []
    with MANIFEST.open("r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            rec = json.loads(line)
            if rec.get("status") == "UNKNOWN":
                unknown.append(rec)

    if not unknown:
        print("No UNKNOWN entries found in manifest.")
        return

    counts = {}
    sizes = {}

    with REPORT.open("w", encoding="utf-8") as out:
        out.write("class\tbytes\tinode\tpath\treason\tetb_blocks\tblocks_output\n")

        for idx, rec in enumerate(unknown, 1):
            result = classify(rec)
            cls = result["class"]
            counts[cls] = counts.get(cls, 0) + 1
            sizes[cls] = sizes.get(cls, 0) + result["size"]

            etb_blocks = ",".join(
                f"ETB{e['level']}:{e['block']}" for e in result["etbs"]
            )

            clean_reason = result["reason"].replace("\t", " ").replace("\n", " ")
            clean_path = result["path"].replace("\t", " ")
            clean_blocks = result["blocks_output"].replace("\t", " ")

            out.write(
                f"{cls}\t{result['size']}\t{result['inode']}\t"
                f"{clean_path}\t{clean_reason}\t{etb_blocks}\t{clean_blocks}\n"
            )
            out.flush()

            print(
                f"[{idx}/{len(unknown)}] {cls:<20} "
                f"{human(result['size']):>10}  {result['path']}"
            )
            if result["etbs"]:
                print(
                    "    "
                    + ", ".join(
                        f"ETB{e['level']}={e['block']}"
                        + (" [DEAD]" if DEAD_START <= e["block"] <= DEAD_END else " [LIVE]")
                        for e in result["etbs"]
                    )
                )

    lines = [
        "PLEX BOOKS UNKNOWN SECONDARY CHECK",
        "==================================",
        f"Unknown entries checked: {len(unknown)}",
        f"Dead block range: {DEAD_START}-{DEAD_END}",
        "",
    ]

    for cls in sorted(counts):
        lines.append(f"{cls:20} {counts[cls]:4d} files  {human(sizes[cls])}")

    lines += [
        "",
        f"Report: {REPORT}",
    ]

    text = "\n".join(lines) + "\n"
    SUMMARY.write_text(text, encoding="utf-8")
    print()
    print(text)

if __name__ == "__main__":
    main()