#!/usr/bin/env python3
from pathlib import Path
import argparse, hashlib, json, shutil, difflib, sys

LAB_ROOT = Path(__file__).resolve().parents[1]
BASELINE_ROOT = LAB_ROOT / "baseline" / "fixture"
WORKING_ROOT = LAB_ROOT / "fixture"
MANIFEST_PATH = LAB_ROOT / "BASELINE_MANIFEST.json"

EXCLUDED_DIRS = {"__pycache__"}
EXCLUDED_SUFFIXES = {".pyc", ".pyo"}

def sha256(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()

def rel_files(root):
    out = {}
    if not root.exists():
        return out
    for p in root.rglob("*"):
        if not p.is_file():
            continue
        if any(part in EXCLUDED_DIRS for part in p.relative_to(root).parts):
            continue
        if p.suffix in EXCLUDED_SUFFIXES:
            continue
        out[p.relative_to(root).as_posix()] = p
    return out

def load_manifest():
    return json.loads(MANIFEST_PATH.read_text(encoding="utf-8"))

def verify_baseline():
    manifest = load_manifest()
    expected = {x["path"]: x["sha256"] for x in manifest["files"]}
    actual = rel_files(BASELINE_ROOT)
    errors = []
    if set(actual) != set(expected):
        errors.append({"kind":"FILE_SET_MISMATCH","expected":sorted(expected),"actual":sorted(actual)})
    for rel, digest in expected.items():
        p = actual.get(rel)
        if p and sha256(p) != digest:
            errors.append({"kind":"HASH_MISMATCH","path":rel,"expected":digest,"actual":sha256(p)})
    result={"command":"verify-baseline","status":"PASS" if not errors else "FAIL","errors":errors}
    print(json.dumps(result,indent=2,sort_keys=True))
    return 0 if not errors else 2

def compute_changes():
    base = rel_files(BASELINE_ROOT); work = rel_files(WORKING_ROOT)
    paths=sorted(set(base)|set(work)); changes=[]
    for rel in paths:
        b=base.get(rel); w=work.get(rel)
        if b is None: changes.append({"path":rel,"status":"ADDED"})
        elif w is None: changes.append({"path":rel,"status":"DELETED"})
        elif sha256(b)!=sha256(w): changes.append({"path":rel,"status":"MODIFIED"})
    return changes,base,work

def cmd_diff():
    if verify_baseline()!=0:
        return 2
    changes,base,work=compute_changes()
    print(json.dumps({"command":"diff","change_count":len(changes),"changes":changes},indent=2,sort_keys=True))
    for item in changes:
        rel=item["path"]; b=base.get(rel); w=work.get(rel)
        if b and w:
            try:
                a=b.read_text(encoding="utf-8").splitlines(True); c=w.read_text(encoding="utf-8").splitlines(True)
            except UnicodeDecodeError:
                continue
            sys.stdout.writelines(difflib.unified_diff(a,c,fromfile=f"baseline/{rel}",tofile=f"fixture/{rel}"))
    return 0

def rollback_check():
    if verify_baseline()!=0:
        return 2
    changes,_,_=compute_changes()
    print(json.dumps({"command":"rollback-check","status":"READY","would_restore":[x["path"] for x in changes]},indent=2,sort_keys=True))
    return 0

def rollback():
    if verify_baseline()!=0:
        return 2
    if WORKING_ROOT.exists():
        shutil.rmtree(WORKING_ROOT)
    shutil.copytree(BASELINE_ROOT, WORKING_ROOT)
    changes,_,_=compute_changes()
    result={"command":"rollback","status":"PASS" if not changes else "FAIL","remaining_changes":changes}
    print(json.dumps(result,indent=2,sort_keys=True))
    return 0 if not changes else 2

def main():
    ap=argparse.ArgumentParser(description="Deterministic LAB-01 baseline/diff/rollback helper")
    ap.add_argument("command",choices=["verify-baseline","diff","rollback-check","rollback"])
    args=ap.parse_args()
    return {"verify-baseline":verify_baseline,"diff":cmd_diff,"rollback-check":rollback_check,"rollback":rollback}[args.command]()

if __name__ == "__main__":
    raise SystemExit(main())