#!/usr/bin/env python3

"""Strict, selective, read-only Plex state salvage from a damaged ext4 LV."""
import argparse
import datetime as dt
import hashlib
import json
import math
import os
import re
import shutil
import stat
import subprocess
import sys
from pathlib import Path

LV = '/dev/mapper/plex--new-plex--new'
DEST_MOUNT = Path('/mnt/recovery')
BASE = '/plexmediaserver/Library/Application Support/Plex Media Server'
DEAD_FIRST = 3418636288
DEAD_LAST = 4395389951
BLOCK_SIZE = 4096
RESERVE = 2 * 1024**3
VERSION = 1

ROOT_FILES = ['Preferences.xml', '.LocalAdminToken']
DATABASE_DIR = f'{BASE}/Plug-in Support/Databases'
PREFERENCES_DIR = f'{BASE}/Plug-in Support/Preferences'
EXTRA_TREES = [
    f'{BASE}/Plug-ins/Hama.bundle',
    f'{BASE}/Plug-ins/Audnexus.bundle',
    f'{BASE}/Scanners/Series',
]
DB_NAME = re.compile(r'^com\.plexapp\.plugins\.library(?:\.blobs)?\.db(?:-(?:[0-9]{4}-[0-9]{2}-[0-9]{2}|wal|shm))?$')
ETB = re.compile(r'\(ETB([0-9]+)\):([0-9]+)')
EXTENT = re.compile(r'\(([0-9]+)(?:-([0-9]+))?\):([0-9]+)(?:-([0-9]+))?')
INODE = re.compile(r'\bInode:\s*([0-9]+)\s+Type:\s*([^\s]+)')
SIZE = re.compile(r'\bSize:\s*([0-9]+)')
OWNER = re.compile(r'\bUser:\s*([0-9]+)\s+Group:\s*([0-9]+)')
MODE = re.compile(r'\bMode:\s*([0-7]{3,4})')

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

def die(message):
    raise SystemExit(f'ERROR: {message}')

def run(cmd):
    return subprocess.run(cmd, capture_output=True, text=True, errors='replace')

def dbg(command):
    env = os.environ.copy()
    env['DEBUGFS_PAGER'] = 'cat'
    return subprocess.run(
        ['debugfs', '-c', '-R', command, LV],
        capture_output=True, text=True, errors='replace', env=env,
    )

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

def warnings(result):
    return '\n'.join(
        line for line in result.stderr.splitlines()
        if line.strip() and not re.match(r'^debugfs [0-9]', line.strip())
    ).strip()

def config(profile):
    prefix = f'plex-state-{profile}'
    return {
        'profile': profile,
        'manifest': DEST_MOUNT / f'{prefix}-manifest.jsonl',
        'summary': DEST_MOUNT / f'{prefix}-scan-summary.txt',
        'recover_log': DEST_MOUNT / f'{prefix}-recovery.jsonl',
        'recover_summary': DEST_MOUNT / f'{prefix}-recovery-summary.txt',
        'tmp': DEST_MOUNT / '.plex-state-tmp',
    }

def preflight(cfg, need_manifest=False):
    if os.geteuid() != 0:
        die('Run as root; the source block device requires read access.')
    for binary in ('debugfs', 'blockdev', 'findmnt'):
        if shutil.which(binary) is None:
            die(f'Required program not installed: {binary}')
    if not Path(LV).exists():
        die(f'LV missing: {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 proceed')
    m = run(['findmnt', '-n', '-o', 'TARGET,SOURCE,FSTYPE', '-T', str(DEST_MOUNT)])
    if m.returncode != 0 or not m.stdout.strip():
        die(f'No mounted recovery filesystem detected at {DEST_MOUNT}')
    parts = m.stdout.split()
    if len(parts) < 3 or parts[0] != str(DEST_MOUNT):
        die(f'{DEST_MOUNT} must itself be the mounted recovery filesystem: {m.stdout.strip()}')
    if LV in m.stdout or 'plex--new' in m.stdout:
        die('Recovery destination appears to be the damaged LV')
    if not os.access(DEST_MOUNT, os.W_OK):
        die(f'{DEST_MOUNT} is not writable')
    if need_manifest and not cfg['manifest'].exists():
        die(f'Missing {cfg["manifest"]}. Scan {cfg["profile"]} first.')
    print(f'Source: {LV} (read-only)')
    print(f'Destination: {m.stdout.strip()}')
    print(f'Dead physical filesystem blocks: {DEAD_FIRST}-{DEAD_LAST}')
    if parts[2].lower() == 'exfat':
        print('NOTICE: destination is exFAT. Tokens and database files are plaintext here;')
        print('        Linux ownership/permissions cannot be preserved on this filesystem.')
    print()

def get_stat(inode_or_path):
    spec = f'<{inode_or_path}>' if isinstance(inode_or_path, int) else quote(inode_or_path)
    r = dbg(f'stat {spec}')
    match = INODE.search(r.stdout)
    if not match:
        return None, (warnings(r) or (r.stdout + '\n' + r.stderr).strip())
    owner = OWNER.search(r.stdout)
    mode = MODE.search(r.stdout)
    size = SIZE.search(r.stdout)
    info = {
        'inode': int(match.group(1)),
        'type': match.group(2),
        'size': int(size.group(1)) if size else None,
        'uid': int(owner.group(1)) if owner else None,
        'gid': int(owner.group(2)) if owner else None,
        'mode': mode.group(1) if mode else None,
        'text': r.stdout,
    }
    return info, warnings(r)

def get_dir_entries(path):
    info, issue = get_stat(path)
    if not info or info['type'] != 'directory':
        return None, f'Cannot open directory: {path} | {issue or info}'
    r = dbg(f'ls -p <{info["inode"]}>')
    issue = warnings(r)
    # Do not silently accept a partial directory listing if debugfs reported errors.
    if issue:
        return None, f'Error listing {path} (inode {info["inode"]}): {issue}'
    entries = []
    for line in r.stdout.splitlines():
        if not line.startswith('/'):
            continue
        bits = line.split('/')
        if len(bits) < 7:
            return None, f'Unparseable directory entry in {path}: {line!r}'
        try:
            inode, mode = int(bits[1]), int(bits[2], 8)
        except ValueError:
            return None, f'Unparseable inode/mode in {path}: {line!r}'
        name = bits[5]
        if not name or name in ('.', '..') or inode == 0:
            continue
        entries.append((name, inode, mode))
    return entries, None

def inspect_file(path, inode):
    info, issue = get_stat(inode)
    rec = {'path': path, 'inode': inode}
    if not info or info['type'] != 'regular' or info['size'] is None:
        return {**rec, 'status': 'UNKNOWN', 'reason': f'Inode unavailable: {issue}'}
    rec.update({key: info[key] for key in ('size', 'uid', 'gid', 'mode')})
    if issue:
        return {**rec, 'status': 'UNKNOWN', 'reason': f'debugfs reported an issue: {issue}'}
    if info['size'] == 0:
        return {**rec, 'status': 'SAFE', 'reason': 'empty file', 'extents': []}
    if 'EXTENTS:' not in info['text']:
        return {**rec, 'status': 'UNKNOWN', 'reason': 'No EXTENTS section (legacy block maps not supported)'}

    data = info['text'].split('EXTENTS:', 1)[1]
    etbs = [{'level': int(l), 'block': int(b)} for l, b in ETB.findall(data)]
    extents = []
    for m in EXTENT.finditer(data):
        ll = int(m.group(1)); lh = int(m.group(2) or m.group(1))
        pl = int(m.group(3)); ph = int(m.group(4) or m.group(3))
        extents.append({'logical_start': ll, 'logical_end': lh, 'physical_start': pl, 'physical_end': ph})
    rec.update({'extents': extents, 'etbs': etbs})

    dead_etbs = [x for x in etbs if DEAD_FIRST <= x['block'] <= DEAD_LAST]
    if dead_etbs:
        return {**rec, 'status': 'ETB_LOST', 'reason': 'Required extent-tree metadata is on dead PV'}
    if not extents:
        return {**rec, 'status': 'ETB_SURVIVES' if etbs else 'UNKNOWN',
                'reason': 'No usable data extents resolved'}

    if any(x['physical_start'] <= DEAD_LAST and x['physical_end'] >= DEAD_FIRST for x in extents):
        return {**rec, 'status': 'LOST', 'reason': 'Data extents overlap dead PV'}

    # Strict check: a file can look SAFE if debugfs returned only the surviving
    # *subset* of extents. Demand complete, non-overlapping, physically plausible
    # block coverage with no holes; intentionally reject sparse/uninitialized files.
    needed = math.ceil(info['size'] / BLOCK_SIZE)
    cursor = 0
    for x in sorted(extents, key=lambda row: row['logical_start']):
        if x['logical_start'] != cursor:
            return {**rec, 'status': 'UNKNOWN', 'reason':
                    f'Incomplete/overlapping logical extent coverage at file block {cursor}'}
        logical_count = x['logical_end'] - x['logical_start'] + 1
        physical_count = x['physical_end'] - x['physical_start'] + 1
        if logical_count != physical_count or logical_count <= 0 or x['physical_start'] == 0:
            return {**rec, 'status': 'UNKNOWN', 'reason': 'Extent length mismatch / invalid physical extent'}
        cursor = x['logical_end'] + 1
    if cursor != needed:
        return {**rec, 'status': 'UNKNOWN', 'reason':
                f'Extent coverage ends at {cursor} blocks; expected {needed}'}

    # Reject unexpected extent syntax (e.g. uninitialized extent annotations)
    # rather than treating any recognized extents as a complete mapping.
    unparsed = ETB.sub('', data)
    unparsed = EXTENT.sub('', unparsed)
    unparsed = re.sub(r'[\s,]+', '', unparsed)
    if unparsed:
        return {**rec, 'status': 'UNKNOWN', 'reason': f'Unexpected EXTENTS text: {unparsed[:100]}'}
    return {**rec, 'status': 'SAFE', 'reason': 'Complete extent coverage outside dead PV'}

def scan(profile):
    cfg = config(profile)
    preflight(cfg)
    selected = {}
    issues = []
    skipped = 0
    visited_dirs = set()

    def select_file(path, inode=None):
        nonlocal skipped
        if path in selected:
            return
        if inode is None:
            info, issue = get_stat(path)
            if not info or info['type'] != 'regular':
                issues.append(f'MISSING/UNREADABLE FILE {path}: {issue or info}')
                return
            inode = info['inode']
        selected[path] = inspect_file(path, inode)
        rec = selected[path]
        if rec['status'] != 'SAFE':
            print(f'{rec["status"]:13} {path}: {rec.get("reason", "")}')

    def select_directory(directory, filter_func=lambda _name, _mode: True, recursive=False):
        nonlocal skipped
        stack = [directory]
        while stack:
            current = stack.pop()
            if current in visited_dirs:
                continue
            visited_dirs.add(current)
            entries, error = get_dir_entries(current)
            if error:
                issues.append(error)
                print('WARNING:', error)
                continue
            for name, inode, mode in entries:
                path = current + '/' + name
                if stat.S_ISDIR(mode):
                    if recursive and filter_func(name, mode):
                        stack.append(path)
                elif stat.S_ISREG(mode):
                    if filter_func(name, mode):
                        select_file(path, inode)
                    else:
                        skipped += 1
                else:
                    issues.append(f'IGNORED nonregular entry: {path} mode={oct(mode)}')

    if profile == 'core':
        for filename in ROOT_FILES:
            select_file(f'{BASE}/{filename}')
        select_directory(DATABASE_DIR, lambda name, _mode: bool(DB_NAME.fullmatch(name)))
        select_directory(PREFERENCES_DIR, lambda name, _mode: name.endswith('.xml'))
    elif profile == 'extras':
        for directory in EXTRA_TREES:
            select_directory(directory, recursive=True)
    else:
        die(f'Unknown profile {profile}')

    counts = {}
    sizes = {}
    for rec in selected.values():
        status = rec['status']
        counts[status] = counts.get(status, 0) + 1
        sizes[status] = sizes.get(status, 0) + (rec.get('size') or 0)

    meta = {
        'version': VERSION, 'profile': profile, 'base': BASE,
        'lv': LV, 'dead_first': DEAD_FIRST, 'dead_last': DEAD_LAST,
        'block_size': BLOCK_SIZE, 'selected_count': len(selected),
        'created_at': dt.datetime.now(dt.timezone.utc).isoformat(),
        'issues': issues,
    }
    new_path = Path(str(cfg['manifest']) + '.new')
    with new_path.open('w', encoding='utf-8') as f:
        f.write(json.dumps({'_meta': meta}, ensure_ascii=False) + '\n')
        for rec in sorted(selected.values(), key=lambda item: item['path']):
            f.write(json.dumps(rec, ensure_ascii=False) + '\n')
    new_path.replace(cfg['manifest'])

    lines = [
        f'PLEX SELECTIVE SCAN: {profile}',
        '=' * 35,
        f'Source: {BASE}',
        f'Selected regular files: {len(selected)}',
        f'Skipped by allowlist: {skipped}',
        f'Scanned directories: {len(visited_dirs)}',
    ]
    for status in ('SAFE', 'LOST', 'ETB_LOST', 'ETB_SURVIVES', 'UNKNOWN'):
        lines.append(f'{status:14} {counts.get(status, 0):>4} files  {human(sizes.get(status, 0)):>12}')
    lines += [
        '',
        f'Available recovery space NOW: {human(shutil.disk_usage(DEST_MOUNT).free)}',
        f'Issues/warnings: {len(issues)}',
    ]
    lines += ['  ' + issue for issue in issues]

    if profile == 'core':
        lines += ['', 'DATABASE CANDIDATES (extent classification; NOT SQLite integrity):']
        for rec in sorted(selected.values(), key=lambda r: r['path']):
            if rec['path'].startswith(DATABASE_DIR + '/'):
                lines.append(f'  {rec["status"]:12} {Path(rec["path"]).name}')
        live = [rec for rec in selected.values() if
                rec['path'].startswith(DATABASE_DIR + '/') and
                Path(rec['path']).name in ('com.plexapp.plugins.library.db', 'com.plexapp.plugins.library.blobs.db')]
        if len(live) < 2:
            lines.append('  WARNING: One or both live database files absent/unreadable')

    lines += ['', f'Manifest: {cfg["manifest"]}',
              'CAUTION: SAFE is based on block mapping; validate recovered databases separately.']
    report = '\n'.join(lines) + '\n'
    cfg['summary'].write_text(report, encoding='utf-8')
    print('\n' + report)
    if issues:
        print('WARNING: scan had missing/inaccessible entries. Review warning lines; do not assume completeness.')

def recover(profile):
    cfg = config(profile)
    preflight(cfg, need_manifest=True)
    with cfg['manifest'].open(encoding='utf-8') as f:
        parsed = [json.loads(line) for line in f if line.strip()]
    if not parsed or '_meta' not in parsed[0]:
        die('No manifest metadata header')
    meta = parsed[0]['_meta']
    if any((meta.get(k) != expected for k, expected in (
        ('profile', profile), ('lv', LV), ('dead_first', DEAD_FIRST),
        ('dead_last', DEAD_LAST), ('block_size', BLOCK_SIZE),
        ('version', VERSION),
    ))):
        die('Manifest geometry/profile/version does not match this script')
    entries = [r for r in parsed[1:] if r.get('status') == 'SAFE']
    print(f'{len(entries)} SAFE files selected, {human(sum(r["size"] for r in entries))} total')
    if meta.get('issues'):
        print('WARNING: scan had missing/unreadable paths; review scan summary.')
    cfg['tmp'].mkdir(exist_ok=True)
    stats = {'RECOVERED': 0, 'ALREADY_PRESENT_UNVERIFIED': 0, 'ERROR': 0, 'NO_SPACE': 0}
    bytes_new = 0
    mode = 'a' if cfg['recover_log'].exists() else 'w'
    with cfg['recover_log'].open(mode, encoding='utf-8') as logfile:
        for idx, rec in enumerate(entries, 1):
            src = rec['path']
            size = rec['size']
            destination = DEST_MOUNT / src.lstrip('/')
            label = None; detail = ''
            digest = None
            # Recheck the live inode's mapping and byte size before copying.
            check = inspect_file(src, rec['inode'])
            if check['status'] != 'SAFE' or check.get('size') != size or check.get('extents') != rec.get('extents'):
                label = 'ERROR'; detail = 'Source inode extent map changed or fails strict safety recheck'
            elif destination.exists():
                if destination.is_file() and destination.stat().st_size == size:
                    label = 'ALREADY_PRESENT_UNVERIFIED'; detail = 'Already present with matching size; not rewritten'
                else:
                    label = 'ERROR'; detail = 'Destination exists with different size or is not a file; refusing overwrite'
            elif shutil.disk_usage(DEST_MOUNT).free < size + RESERVE:
                label = 'NO_SPACE'; detail = 'Need file size plus 2 GiB free-space reserve'
            else:
                tmp = cfg['tmp'] / f'{profile}-{rec["inode"]}.part'
                if tmp.exists():
                    tmp.unlink()
                result = dbg(f'dump <{rec["inode"]}> {tmp}')
                issue = warnings(result)
                if issue:
                    label = 'ERROR'; detail = f'debugfs reported: {issue[:400]}'
                elif not tmp.exists() or tmp.stat().st_size != size:
                    label = 'ERROR'; detail = 'Missing output file or byte-size mismatch'
                else:
                    digest_obj = hashlib.sha256()
                    with tmp.open('rb') as source:
                        for chunk in iter(lambda: source.read(4 * 1024 * 1024), b''):
                            digest_obj.update(chunk)
                    digest = digest_obj.hexdigest()
                    destination.parent.mkdir(parents=True, exist_ok=True)
                    try:
                        tmp.replace(destination)
                        label = 'RECOVERED'; bytes_new += size
                    except OSError as err:
                        label = 'ERROR'; detail = f'Rename failed: {err}'
                if label == 'ERROR' and tmp.exists():
                    try:
                        tmp.unlink()
                    except OSError:
                        pass
            if label in ('ERROR', 'NO_SPACE'):
                print(f'{label:10} {src}: {detail}')
            else:
                print(f'{label:26} [{idx}/{len(entries)}] {src}')
            stats[label] += 1
            logfile.write(json.dumps({
                'status': label, 'path': src, 'size': size, 'inode': rec['inode'],
                'sha256': digest, 'detail': detail,
            }, ensure_ascii=False) + '\n')
            logfile.flush()
            if label == 'NO_SPACE':
                break

    lines = [
        f'PLEX SELECTIVE RECOVERY: {profile}', '=' * 35,
        f'Newly recovered: {stats["RECOVERED"]}, {human(bytes_new)}',
        f'Already present (size-only): {stats["ALREADY_PRESENT_UNVERIFIED"]}',
        f'Errors: {stats["ERROR"]}', f'Space stops: {stats["NO_SPACE"]}',
        f'Remaining free: {human(shutil.disk_usage(DEST_MOUNT).free)}',
        f'Manifest: {cfg["manifest"]}',
        f'Hash log: {cfg["recover_log"]}',
        'Hashes verify extracted copies, NOT correctness of Plex databases.',
    ]
    report = '\n'.join(lines) + '\n'
    cfg['recover_summary'].write_text(report, encoding='utf-8')
    print('\n' + report)
    return 1 if stats['ERROR'] or stats['NO_SPACE'] else 0

def main():
    os.umask(0o077)
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('action', choices=('scan', 'recover'))
    parser.add_argument('profile', choices=('core', 'extras'), nargs='?', default='core')
    args = parser.parse_args()
    if args.action == 'scan':
        scan(args.profile)
    else:
        sys.exit(recover(args.profile))

if __name__ == '__main__':
    main()