import hashlib
import os
import stat
import sys

def fail(msg):
    print(msg, file=sys.stderr)
    sys.exit(1)

def symlink_resolves_inside(root, link_path, link_text):
    parent = os.path.dirname(os.path.abspath(link_path))
    resolved = os.path.normpath(os.path.join(parent, link_text))
    root = os.path.abspath(root)
    return resolved == root or resolved.startswith(root + os.sep)

def file_sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as fh:
        while True:
            chunk = fh.read(1024 * 1024)
            if not chunk:
                return digest.hexdigest()
            digest.update(chunk)

def validate(quarantine_root, target_roots, legacy):
    qroot = os.path.abspath(quarantine_root)
    troots = [os.path.abspath(root) for root in target_roots]
    if not os.path.isdir(qroot):
        fail("subset:not-a-directory")
    if not troots or any(not os.path.exists(root) for root in troots):
        fail("subset:target-missing")
    for dirpath, dirnames, filenames in os.walk(qroot, topdown=True, followlinks=False):
        rel_dir = os.path.relpath(dirpath, qroot)
        if rel_dir == ".":
            rel_dir = ""
        descend = []
        for name in sorted(set(dirnames) | set(filenames)):
            # Python bytecode caches regenerate wherever a shipped script runs;
            # they are derived state, never payload, and never block takeover.
            if name == "__pycache__" or name.endswith(".pyc"):
                continue
            rel = name if not rel_dir else f"{rel_dir}/{name}"
            qpath = os.path.join(dirpath, name)
            try:
                qst = os.lstat(qpath)
            except OSError:
                fail(f"subset:unreadable:{rel}")
            targets = []
            target_error = False
            for troot in troots:
                tpath = os.path.join(troot, rel)
                try:
                    targets.append((tpath, os.lstat(tpath)))
                except FileNotFoundError:
                    pass
                except OSError:
                    target_error = True
            if not targets:
                fail(f"subset:target-unreadable:{rel}" if target_error else f"subset:unknown:{rel}")
            qmode = stat.S_IFMT(qst.st_mode)
            typed = [(path, tst) for path, tst in targets if stat.S_IFMT(tst.st_mode) == qmode]
            if not typed:
                fail(f"subset:type-mismatch:{rel}")
            if qmode == stat.S_IFREG:
                matched = False
                for tpath, tst in typed:
                    if qst.st_size != tst.st_size:
                        continue
                    with open(qpath, "rb") as qf, open(tpath, "rb") as tf:
                        while True:
                            qb = qf.read(1024 * 1024)
                            tb = tf.read(1024 * 1024)
                            if qb != tb:
                                break
                            if not qb:
                                matched = True
                                break
                    if matched:
                        break
                if not matched and file_sha256(qpath) in legacy.get(rel, set()):
                    matched = True
                if not matched:
                    fail(f"subset:content-mismatch:{rel}")
            elif qmode == stat.S_IFLNK:
                try:
                    qlink = os.readlink(qpath)
                except OSError:
                    fail(f"subset:symlink-unreadable:{rel}")
                if not symlink_resolves_inside(qroot, qpath, qlink):
                    fail(f"subset:unsafe-symlink:{rel}")
                matched = False
                for tpath, _tst in typed:
                    try:
                        if qlink == os.readlink(tpath):
                            matched = True
                            break
                    except OSError:
                        continue
                if not matched:
                    fail(f"subset:symlink-mismatch:{rel}")
            elif qmode == stat.S_IFDIR:
                descend.append(name)
            else:
                fail(f"subset:special-file:{rel}")
        dirnames[:] = descend

if __name__ == "__main__":
    if len(sys.argv) < 3:
        fail("subset:usage")
    args = sys.argv[2:]
    legacy = {}
    if "--legacy" in args:
        split = args.index("--legacy")
        target_roots, specs = args[:split], args[split + 1:]
        for spec in specs:
            try:
                rel, digest = spec.rsplit("=", 1)
            except ValueError:
                fail("subset:legacy-usage")
            if not rel or os.path.isabs(rel) or ".." in rel or len(digest) != 64 or any(c not in "0123456789abcdef" for c in digest):
                fail("subset:legacy-usage")
            legacy.setdefault(rel, set()).add(digest)
    else:
        target_roots = args
    validate(sys.argv[1], target_roots, legacy)