#!/usr/bin/env python3
"""Recursively extract every archive under the source roots into one tree.

    python3 tools/extract_all.py <src_root>... --out <extracted_root>

Behavior:
- Every archive (zip/7z/rar/tar/gz/bz2/xz) is extracted into a directory
  named after the archive (suffix stripped), preserving the relative path
  under the source root:  <src>/zips/The Twin Lance.7z
      -> <out>/1/zips/The Twin Lance/...
- Archives found *inside* an extraction are extracted too (BFS, until no
  archive remains), then deleted - the extracted tree holds no archives.
- Plain files are hardlinked into the tree (fallback: copy), so nothing
  is lost and no data is duplicated.
- The source roots are never modified (originals stay).
- A provenance manifest `<out>/extract_manifest.jsonl` records one line
  per extraction: {"archive": ..., "dest": ..., "unit": ...} for the
  sorter's --provenance unit mode.
- Failures (bad download/HTML page, wrong password, disk full) are moved
  to `<out>/.broken/<relpath>` and logged; the run never stops.

Re-runs are safe: destinations recorded in the manifest are skipped (a
partial dir from a crash is NOT skipped - it gets re-extracted); a final
sweep expands archives stranded by interrupted nested expansions.
"""

from __future__ import annotations

import argparse
import logging
import os
import shutil
import subprocess
import sys
import tempfile
from collections import deque
from pathlib import Path

log = logging.getLogger("extract_all")

ARCHIVE_SUFFIXES = {".zip", ".7z", ".rar", ".tar", ".gz", ".bz2", ".xz"}
EXTRACT_TIMEOUT = 4 * 3600  # seconds per archive
MIN_FREE_BYTES = 20 * 1024**3  # stop extracting below this free space
RUN_ARTIFACTS = {
    "export.json", "export.html", "state.json", "manifest.jsonl",
    "failed_urls.txt", "catalogue.md", "catalogue.jsonl",
    "sorter_applied.jsonl", "catalogue_moves.jsonl", "urls.txt",
    "urls_0807.txt", "urls_gofile.txt", "failed_0807.txt",
    "extract_manifest.jsonl",
}


def archive_suffixes(name: str) -> list[str]:
    """All trailing archive suffixes, e.g. 'x.tar.gz' -> ['.gz', '.tar']."""
    out = []
    p = Path(name)
    while p.suffix.lower() in ARCHIVE_SUFFIXES:
        out.append(p.suffix.lower())
        p = p.with_suffix("")
    return out


def is_archive(path: Path) -> bool:
    return bool(archive_suffixes(path.name))


def display_unit(name: str) -> str:
    """Archive name with all archive suffixes stripped: the unit folder."""
    p = Path(name)
    while p.suffix.lower() in ARCHIVE_SUFFIXES:
        p = p.with_suffix("")
    return p.name


def free_bytes(path: Path) -> int:
    return shutil.disk_usage(path).free


def _fix_component(s: str) -> str:
    """Repair filenames 7z wrote with raw legacy bytes (cp437 zip entries).

    Raw non-UTF8 bytes in a name cannot be created on the CIFS share
    (observed: errno 2 for a name containing byte 0x87). Decode the raw
    bytes as cp437 so e.g. a raw-0x87 name becomes a proper 'ç' name.
    """
    try:
        s.encode("utf-8")
        return s
    except UnicodeEncodeError:
        pass
    raw = os.fsencode(s)
    try:
        return raw.decode("utf-8")
    except UnicodeDecodeError:
        return raw.decode("cp437", "replace")


def _move_tree(tmp: Path, dest: Path) -> int:
    """Move a freshly extracted local tree onto the (possibly CIFS) dest.
    Returns the number of files that could not be moved."""
    bad = 0
    for f in sorted(tmp.rglob("*")):
        rel = f.relative_to(tmp)
        tgt = dest.joinpath(*(_fix_component(p) for p in rel.parts))
        if f.is_dir():
            tgt.mkdir(parents=True, exist_ok=True)
            continue
        tgt.parent.mkdir(parents=True, exist_ok=True)
        try:
            if tgt.exists() and tgt.is_file() \
                    and tgt.stat().st_size == f.stat().st_size:
                f.unlink()
                continue
            if tgt.exists():
                tgt.unlink()
            shutil.move(str(f), str(tgt))
        except OSError as exc:
            bad += 1
            log.error("Cannot place %s: %s", tgt, exc)
    return bad


def _run_7z(cmd: list[str]) -> tuple[int, str]:
    try:
        # 7z prints legacy-encoded inner filenames (non-UTF8 bytes); decode
        # tolerantly or the whole run crashes (seen with byte 0x87).
        proc = subprocess.run(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
            encoding="utf-8", errors="replace", timeout=EXTRACT_TIMEOUT,
        )
    except subprocess.TimeoutExpired:
        return -1, f"extraction timed out after {EXTRACT_TIMEOUT}s"
    except OSError as exc:
        return -1, f"7z failed to run: {exc}"
    tail = "\n".join((proc.stdout or "").splitlines()[-5:])
    return proc.returncode, tail


def extract(archive: Path, dest: Path) -> tuple[bool, str, Path]:
    """7z x into dest. Empty -p avoids interactive prompts (no password).

    Returns (ok, error, dest_used): on a name collision (the archive stem
    exists as a *file*, e.g. 'model.stl.zip' next to 'model.stl') the
    destination gets a '_dup' suffix instead of crashing.

    If 7z cannot write directly onto the (CIFS) destination - legacy
    cp437 entry names - the archive is extracted to a local temp dir and
    moved over with repaired filenames.
    """
    actual = dest
    for _ in range(5):
        try:
            actual.mkdir(parents=True, exist_ok=True)
            break
        except FileExistsError:
            actual = dest.with_name(dest.name + "_dup")
        except NotADirectoryError as exc:
            return False, f"cannot create destination {actual}: {exc}", dest
    else:
        return False, f"cannot create destination {dest} (repeated collisions)", dest
    rc, tail = _run_7z(["7z", "x", "-y", "-p", f"-o{actual}", str(archive)])
    if rc in (0, 1):  # 0 ok, 1 warnings (e.g. minor CRC)
        if rc == 1:
            log.warning("7z reported warnings extracting %s", archive.name)
        return True, "", actual
    if rc == -1:
        return False, tail, actual
    # Fallback: local extraction, then move with repaired filenames.
    tmp = Path(tempfile.mkdtemp(prefix="7zext-", dir="/var/tmp"))
    try:
        rc2, tail2 = _run_7z(["7z", "x", "-y", "-p", f"-o{tmp}", str(archive)])
        if rc2 not in (0, 1):
            return False, f"7z exit {rc2}: {tail2}", actual
        bad = _move_tree(tmp, actual)
        if bad:
            return False, f"{bad} files could not be moved onto the share", actual
        return True, "", actual
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def move_to_broken(path: Path, out_root: Path, src_root: Path, note: str) -> Path:
    rel = path.relative_to(src_root)
    dest = out_root / ".broken" / rel
    dest.parent.mkdir(parents=True, exist_ok=True)
    if dest.exists():
        dest = dest.with_name(dest.name + ".1")
    return shutil.move(str(path), str(dest))


def is_macosx_junk(p: Path) -> bool:
    """AppleDouble resource-fork 'zips' (.__MACOSX/._x.zip) - no content."""
    return "__MACOSX" in p.parts or p.name.startswith("._")


def expand_nested(dest: Path, out_root: Path, manifest, src_root: Path,
                  done: set[str]) -> int:
    """Extract archives inside `dest` until none remain. Returns count."""
    done_count = 0
    queue: deque[Path] = deque([dest])
    while queue:
        cur = queue.popleft()
        for p in sorted(cur.rglob("*")):
            if not p.is_file() or not is_archive(p):
                continue
            if is_macosx_junk(p):
                p.unlink()
                log.info("Junk (macOS AppleDouble) removed: %s", p.name)
                continue
            stem_dir = p.with_name(display_unit(p.name))
            if str(stem_dir.relative_to(out_root)) in done:
                # already extracted in a previous run; drop the archive copy
                p.unlink()
                log.info("Nested %s already extracted - archive removed", p.name)
                continue
            if free_bytes(out_root) < MIN_FREE_BYTES:
                log.error("Out of disk space - skipping %s (kept in place)", p)
                continue
            ok, err, actual = extract(p, stem_dir)
            if not ok:
                log.error("FAILED nested extraction %s: %s", p, err)
                move_to_broken(p, out_root, out_root, err)
                continue
            manifest.write(f"{json_line(p.relative_to(out_root), actual)}\n")
            done.add(str(actual.relative_to(out_root)))
            p.unlink()  # archive copy consumed
            done_count += 1
            queue.append(actual)
    return done_count


def json_line(archive_rel: Path, dest: Path) -> str:
    import json
    return json.dumps({
        "archive": str(archive_rel),
        "dest": str(dest.relative_to(_OUT_ROOT)),
        "unit": display_unit(archive_rel.name),
    }, ensure_ascii=False)


_OUT_ROOT: Path = Path(".")  # set in main; used by json_line


def process_root(src_root: Path, out_root: Path, manifest, done: set[str]) -> tuple[int, int, int]:
    """Returns (extracted, mirrored, failed)."""
    idx = src_root.name
    extracted = mirrored = failed = 0
    for f in sorted(src_root.rglob("*")):
        if not f.is_file():
            continue
        if (f.name in RUN_ARTIFACTS or f.name.startswith(".")
                or f.name.endswith(".part")):
            continue
        rel = f.relative_to(src_root)
        if is_archive(f):
            dest_rel = Path(idx) / Path(str(rel.with_name(display_unit(f.name))))
            dest = out_root / dest_rel
            # Skip on a manifest record, never on "dir non-empty": a crash
            # mid-extraction leaves a partial dir that must be redone.
            if str(dest_rel) in done:
                log.info("Skip (already extracted): %s", rel)
                continue
            if free_bytes(out_root) < MIN_FREE_BYTES:
                log.error("Out of disk space - skipping archive %s", f)
                failed += 1
                continue
            ok, err, actual = extract(f, dest)
            if not ok:
                log.error("FAILED %s: %s", rel, err)
                move_to_broken(f, out_root, src_root, err)
                failed += 1
                continue
            manifest.write(f"{json_line(Path(idx) / rel, actual)}\n")
            done.add(str(actual.relative_to(out_root)))
            extracted += 1
            log.info("Extracted %s -> %s", rel, actual.relative_to(out_root))
            expand_nested(actual, out_root, manifest, src_root, done)
        else:
            link_dest = out_root / idx / rel
            if link_dest.exists():
                continue
            link_dest.parent.mkdir(parents=True, exist_ok=True)
            try:
                os.link(f, link_dest)
            except OSError:
                shutil.copy2(f, link_dest)
            mirrored += 1
    return extracted, mirrored, failed


def load_done(manifest_path: Path) -> set[str]:
    """Dest dirs already recorded in the manifest (relative to out root)."""
    import json
    done: set[str] = set()
    if manifest_path.exists():
        with manifest_path.open(encoding="utf-8") as fh:
            for line in fh:
                line = line.strip()
                if not line:
                    continue
                try:
                    done.add(json.loads(line)["dest"])
                except (json.JSONDecodeError, KeyError):
                    continue
    return done


def sweep(out_root: Path, manifest, done: set[str]) -> int:
    """Expand any archive left anywhere in the tree (interrupted nested
    expansions leave stranded copies even when the parent has a record)."""
    import json
    swept = 0
    for p in sorted(out_root.rglob("*")):
        if not p.is_file() or not is_archive(p):
            continue
        if ".broken" in p.relative_to(out_root).parts:
            continue
        if is_macosx_junk(p):
            p.unlink()
            log.info("Junk (macOS AppleDouble) removed: %s", p.name)
            continue
        stem_dir = p.with_name(display_unit(p.name))
        if str(stem_dir.relative_to(out_root)) in done:
            p.unlink()
            log.info("Sweep: %s already extracted - archive removed", p.name)
            continue
        if free_bytes(out_root) < MIN_FREE_BYTES:
            log.error("Out of disk space - sweep skipping %s", p)
            continue
        ok, err, actual = extract(p, stem_dir)
        if not ok:
            log.error("FAILED sweep extraction %s: %s", p, err)
            move_to_broken(p, out_root, out_root, err)
            continue
        manifest.write(f"{json_line(p.relative_to(out_root), actual)}\n")
        done.add(str(actual.relative_to(out_root)))
        p.unlink()
        swept += 1
        expand_nested(actual, out_root, manifest, out_root, done)
    return swept


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("sources", nargs="+", type=Path)
    parser.add_argument("--out", type=Path, required=True, metavar="DIR")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO,
                        format="%(asctime)s %(levelname)-7s %(message)s",
                        datefmt="%H:%M:%S")

    global _OUT_ROOT
    _OUT_ROOT = args.out
    args.out.mkdir(parents=True, exist_ok=True)

    totals = [0, 0, 0]
    done = load_done(args.out / "extract_manifest.jsonl")
    log.info("%d destinations already recorded in manifest", len(done))
    with (args.out / "extract_manifest.jsonl").open("a", encoding="utf-8") as manifest:
        for src in args.sources:
            log.info("== processing %s", src)
            counts = process_root(src, args.out, manifest, done)
            log.info("== %s: %d extracted, %d mirrored, %d failed",
                     src.name, *counts)
            for i, c in enumerate(counts):
                totals[i] += c
        swept = sweep(args.out, manifest, done)
    log.info("TOTAL: %d archives extracted, %d files mirrored, %d failed, %d swept",
             *totals, swept)
    if totals[2]:
        log.warning("Failures were moved under %s - inspect and re-run", args.out / ".broken")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(130)