#!/usr/bin/env python3
"""Fail-closed preservation and retirement for Git worktrees."""
import fcntl, json, os, re, shutil, subprocess, sys, tempfile
from pathlib import Path

SECRET = re.compile(r'(^|/)(\.env(?:\.|$)|.*\.(pem|key|p12|pfx|kdbx)|credentials?\.?[^/]*|secrets?\.?[^/]*)$', re.I)
GENERATED = re.compile(r'(^|/)(node_modules|\.git|\.cache|cache|dist|build|coverage|\.pytest_cache|\.mypy_cache|\.turbo|logs?)(/|$)|\.(log|tmp|tsbuildinfo)$', re.I)

class Refusal(Exception): pass

def run(repo, *args, env=None, check=True):
    p = subprocess.run(["git", "-C", str(repo), *args], text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env)
    if check and p.returncode: raise Refusal(p.stderr.strip() or f"git {' '.join(args)} failed")
    return p

def die(msg):
    print(f"archive-worktrees: REFUSED: {msg}", file=sys.stderr); return 2

def root_for(path):
    try:
        top=Path(subprocess.check_output(["git","-C",str(path),"rev-parse","--show-toplevel"], text=True).strip()).resolve()
        raw=Path(subprocess.check_output(["git","-C",str(top),"rev-parse","--git-common-dir"], text=True).strip())
        common=(raw if raw.is_absolute() else top/raw).resolve()
        return (common.parent if common.name==".git" else common).resolve()
    except Exception: raise Refusal("not inside a Git worktree")

def worktree_info(repo):
    lines=run(repo,"worktree","list","--porcelain").stdout.splitlines(); out=[]; cur=None
    for line in lines:
        if line.startswith("worktree "):
            cur={"path":Path(line[9:]).resolve()}
            out.append(cur)
        elif cur and line.startswith("branch "): cur["branch"]=line[7:]
        elif cur and line=="detached": cur["branch"]=None
    return out

def target(repo, arg=None):
    lexical=Path(arg).expanduser() if arg else repo
    if arg is not None and path_has_symlink(lexical): raise Refusal("target path contains a symlink")
    p=lexical.resolve()
    infos=worktree_info(repo)
    match=next((x for x in infos if x["path"]==p),None)
    if not match: raise Refusal("target is not a registered worktree")
    common=Path(run(repo,"rev-parse","--git-common-dir").stdout.strip())
    common=(common if common.is_absolute() else repo/common).resolve()
    main=(common.parent if common.name==".git" else common).resolve()
    default=run(repo,"symbolic-ref","--quiet","--short","refs/remotes/origin/HEAD",check=False).stdout.strip()
    default=default.removeprefix("origin/") or "main"
    if p==main: raise Refusal("refusing shared/default checkout")
    if match.get("branch") == f"refs/heads/{default}": raise Refusal("refusing shared/default checkout")
    if match.get("branch") is None: raise Refusal("refusing detached worktree")
    return p, match

def active(path):
    ancestors={str(os.getpid())}
    current=str(os.getpid())
    while current.isdigit():
        try: current=(Path("/proc")/current/"stat").read_text().split()[3]
        except (OSError,IndexError,ValueError): break
        if current in ancestors: break
        ancestors.add(current)
    for proc in Path("/proc").glob("[0-9]*"):
        if proc.name in ancestors: continue
        try:
            cwd=(proc/"cwd").resolve()
            if cwd==path or path in cwd.parents: return proc.name
            if str(path).encode() in (proc/"cmdline").read_bytes(): return proc.name
        except (OSError,RuntimeError): pass
    return None

def path_has_symlink(path):
    current=Path(path.anchor) if path.is_absolute() else Path.cwd()
    parts=path.parts[1:] if path.is_absolute() else path.parts
    for part in parts:
        current/=part
        if current.is_symlink(): return True
    return False

def inventory(path, strict=False):
    records=run(path,"status","--porcelain=v1","--untracked-files=all","--ignored=matching","-z").stdout.split("\0")
    entries=[]; uncertain=[]
    source_ext=re.compile(r'\.(?:py|js|jsx|ts|tsx|mjs|cjs|rs|go|java|rb|php|md|txt|json|ya?ml|toml|ini|cfg|conf|sh|css|scss|html|sql|lock)$',re.I)
    i=0
    while i < len(records):
        record=records[i]; i += 1
        if len(record)<3: continue
        code,name=record[:2],record[3:]
        rename_source=None
        if "R" in code or "C" in code:
            if i >= len(records): raise Refusal("malformed NUL status rename record")
            rename_source=records[i]; i += 1
        secret=bool(SECRET.search(name) or (rename_source and SECRET.search(rename_source)))
        generated=bool(GENERATED.search(name) or (rename_source and GENERATED.search(rename_source))); ignored=code=="!!"
        if secret:
            entries.append({"path":name,"kind":"ignored" if ignored else "secret","excluded":True})
            if strict: uncertain.append(("secret path",name))
        elif generated:
            entries.append({"path":name,"kind":"ignored" if ignored else "generated","excluded":True})
        elif ignored:
            entries.append({"path":name,"kind":"ignored","excluded":True})
            uncertain.append(("uncertain ignored path",name))
        elif code=="??":
            full=path/name
            if full.is_symlink():
                entries.append({"path":name,"kind":"untracked","excluded":True}); uncertain.append(("symlink",name))
            elif source_ext.search(name):
                entries.append({"path":name,"kind":"untracked","excluded":False})
            else:
                entries.append({"path":name,"kind":"untracked","excluded":True}); uncertain.append(("uncertain untracked path",name))
        else:
            kind="staged" if code[0] not in " ?" else "unstaged"
            entries.append({"path":name,"kind":kind,"excluded":False})
    if strict and uncertain:
        reason,name=uncertain[0]; raise Refusal(f"{reason}: {name}")
    return entries

def archive_dir(repo):
    d=Path(run(repo,"rev-parse","--git-path","worktree-archives").stdout.strip())
    if not d.is_absolute(): d=repo/d
    d.mkdir(parents=True,exist_ok=True); return d

def manifest_for(repo,path): return archive_dir(repo)/(path.name+".json")

def manifest_for_branch(repo, branch):
    matches=[]
    for candidate in archive_dir(repo).glob("*.json"):
        try: data=json.loads(candidate.read_text())
        except (OSError,ValueError): continue
        if data.get("archive_branch") == branch:
            matches.append((candidate,data))
    if len(matches) != 1:
        raise Refusal("preservation receipt for archive branch is missing or ambiguous")
    return matches[0]

def remote_branch_sha(repo, branch):
    result=run(repo,"ls-remote","--refs","origin",f"refs/heads/{branch}",check=False)
    if result.returncode:
        raise Refusal("canonical remote archive state unavailable")
    fields=result.stdout.split()
    if not fields: return None
    if not re.fullmatch(r"[0-9a-f]{40}", fields[0]):
        raise Refusal("canonical remote archive SHA invalid")
    return fields[0]

def commit_tree(repo, ref):
    result=run(repo,"rev-parse","--verify",f"{ref}^{{tree}}",check=False)
    return result.stdout.strip() if result.returncode == 0 else None

def commit_parents(repo, ref):
    result=run(repo,"rev-list","--parents","-n","1",ref,check=False)
    if result.returncode: return None
    fields=result.stdout.split()
    return fields[1:] if fields else None

def validate_archive_commit(repo, ref, head, snapshot_tree, index_tree, head_tree):
    if commit_tree(repo, ref) != snapshot_tree:
        raise Refusal("existing archive branch snapshot mismatch")
    parents=commit_parents(repo, ref)
    if not parents or parents[0] != head:
        raise Refusal("existing archive branch parent mismatch")
    if index_tree != head_tree:
        if len(parents) != 2 or commit_tree(repo, parents[1]) != index_tree:
            raise Refusal("existing archive branch index snapshot mismatch")
    elif len(parents) != 1:
        raise Refusal("existing archive branch parent mismatch")

def preserve(repo, path):
    lock_path=archive_dir(repo)/(path.name+".lock")
    lock=lock_path.open("w")
    fcntl.flock(lock,fcntl.LOCK_EX)
    if active(path): raise Refusal("target has active processes")
    entries=inventory(path, strict=True)
    tracked=run(path,"ls-files","-z").stdout.split("\0")
    entries.extend({"path":name,"kind":"tracked-excluded","excluded":True} for name in tracked if name and (SECRET.search(name) or GENERATED.search(name)))
    branch=run(path,"symbolic-ref","--short","HEAD").stdout.strip(); head=run(path,"rev-parse","HEAD").stdout.strip()
    remote=run(path,"remote","get-url","origin",check=False).stdout.strip()
    if not remote: raise Refusal("canonical remote origin unavailable")
    archive_branch=f"archive/{branch.replace('/','-')}-{head[:12]}"
    ad=archive_dir(repo); manifest=manifest_for(repo,path)
    if manifest.is_file(): raise Refusal("preservation receipt already exists")
    status_by_path={}
    for entry in entries:
        status=run(path,"status","--porcelain=v1","--",entry["path"]).stdout[:2]
        if status: status_by_path[entry["path"]]=status
    status_snapshot=run(path,"status","--porcelain=v1","--untracked-files=all","--ignored=matching","-z").stdout
    original_index_tree=run(path,"write-tree").stdout.strip()
    head_tree=commit_tree(path,head)
    if not head_tree: raise Refusal("original HEAD tree unavailable")
    with tempfile.TemporaryDirectory() as td:
        idx=Path(td)/"index"; env=os.environ.copy(); env["GIT_INDEX_FILE"]=str(idx)
        run(path,"read-tree",head,env=env); run(path,"add","-A",env=env)
        post_add=run(path,"status","--porcelain=v1","--untracked-files=all","--ignored=matching","-z").stdout
        if post_add != status_snapshot:
            raise Refusal("worktree changed during preservation")
        safe=[e["path"] for e in entries if not e["excluded"]]
        allpaths=[e["path"] for e in entries]
        drop=[x for x in allpaths if x not in safe]
        if drop: run(path,"reset","--",*drop,env=env)
        snapshot_tree=run(path,"write-tree",env=env).stdout.strip()
    local_sha=run(path,"rev-parse","--verify",f"refs/heads/{archive_branch}",check=False).stdout.strip()
    remote_sha=remote_branch_sha(path,archive_branch)
    if local_sha and remote_sha and local_sha != remote_sha:
        raise Refusal("remote SHA mismatch: local and canonical archive branches differ")
    if not local_sha and remote_sha:
        run(path,"fetch","origin",f"refs/heads/{archive_branch}:refs/heads/{archive_branch}")
        local_sha=run(path,"rev-parse","--verify",f"refs/heads/{archive_branch}").stdout.strip()
        if local_sha != remote_sha:
            raise Refusal("remote SHA mismatch: fetched archive branch differs")
    if local_sha:
        validate_archive_commit(path,f"refs/heads/{archive_branch}",head,snapshot_tree,original_index_tree,head_tree)
    else:
        commit_env={**os.environ,"GIT_AUTHOR_NAME":"archive-worktrees","GIT_AUTHOR_EMAIL":"archive-worktrees@localhost","GIT_COMMITTER_NAME":"archive-worktrees","GIT_COMMITTER_EMAIL":"archive-worktrees@localhost"}
        index_commit=None
        if original_index_tree != head_tree:
            index_commit=run(path,"commit-tree",original_index_tree,"-p",head,env=commit_env).stdout.strip()
        parents=["-p",head]
        if index_commit: parents.extend(["-p",index_commit])
        commit=run(path,"commit-tree",snapshot_tree,*parents,env=commit_env).stdout.strip()
        run(path,"branch",archive_branch,commit)
        local_sha=commit
    if remote_sha != local_sha:
        run(path,"push","origin",f"{archive_branch}:{archive_branch}")
    remote_sha=remote_branch_sha(path,archive_branch)
    if remote_sha != local_sha: raise Refusal("remote SHA mismatch")
    parents=commit_parents(path,f"refs/heads/{archive_branch}") or []
    index_commit=parents[1] if original_index_tree != head_tree and len(parents) == 2 else None
    staged=ad/(path.name+".staged.json")
    staged_entries=[]
    for entry in entries:
        code=status_by_path.get(entry["path"],"")
        if code and code[0] not in " ?":
            staged_entries.append({"path":entry["path"],"index_status":code[0],"worktree_status":code[1]})
    staged_data={"version":2,"original_head":head,"index_tree":original_index_tree,"staged_paths":[e["path"] for e in staged_entries],"staged_entries":staged_entries}
    staged.write_text(json.dumps(staged_data,sort_keys=True)+"\n")
    data={"version":2,"original_path":str(path),"original_branch":branch,"original_head":head,"archive_branch":archive_branch,"snapshot_sha":local_sha,"remote_sha":remote_sha,"remote":remote,"index_tree":original_index_tree,"index_commit":index_commit,"exclusions":[e["path"] for e in entries if e["excluded"]],"manifest_path":str(manifest),"staged_metadata_path":str(staged),"restore_command":f"archive-worktrees restore {archive_branch} <path>"}
    manifest.write_text(json.dumps(data,indent=2)+"\n")
    print(json.dumps(data,sort_keys=True)); return 0

def retire(repo,path):
    if active(path): raise Refusal("target has active processes")
    _, info=target(repo,str(path)); m=manifest_for(repo,path)
    if not m.is_file(): raise Refusal("preservation receipt missing")
    try: data=json.loads(m.read_text())
    except Exception: raise Refusal("preservation receipt invalid")
    if data.get("original_path")!=str(path) or data.get("original_branch")!=info.get("branch","").removeprefix("refs/heads/"): raise Refusal("preservation receipt target mismatch")
    current_remote=run(repo,"remote","get-url","origin",check=False).stdout.strip()
    if not current_remote or current_remote != data.get("remote"):
        raise Refusal("canonical remote origin changed since preservation")
    remote=run(repo,"ls-remote","origin",f"refs/heads/{data.get('archive_branch','')}",check=False).stdout.split()
    if not remote or remote[0]!=data.get("remote_sha") or data.get("snapshot_sha")!=data.get("remote_sha"): raise Refusal("preservation receipt is stale or archival remote SHA no longer verifies")
    run(repo,"worktree","remove",str(path)); print(f"retired {path}"); return 0

def restore(repo, branch, dest):
    if any(part==".." for part in Path(dest).parts): raise Refusal("path traversal refused")
    d=Path(dest).expanduser().resolve(); root=repo.resolve(); worktrees=(root/".worktrees").resolve()
    inside=root in d.parents
    if d==root or (inside and d.parent!=worktrees): raise Refusal("restore path inside repository must be a direct .worktrees child")
    if d.exists() or d.is_symlink(): raise Refusal("restore path already exists")
    if not re.fullmatch(r"archive/[A-Za-z0-9._/-]+",branch): raise Refusal("invalid archive branch")
    remote=run(repo,"ls-remote","--refs","origin",f"refs/heads/{branch}",check=False).stdout.split()
    if len(remote)<2: raise Refusal("archive branch not found on canonical remote")
    remote_sha=remote[0]
    local=run(repo,"rev-parse","--verify",f"refs/heads/{branch}",check=False).stdout.strip()
    if not local:
        run(repo,"fetch","origin",f"refs/heads/{branch}:refs/heads/{branch}")
    elif local!=remote_sha:
        raise Refusal("local archive branch differs from canonical remote")
    if run(repo,"rev-parse","--verify",f"refs/heads/{branch}").stdout.strip()!=remote_sha: raise Refusal("fetched archive SHA mismatch")
    _, data=manifest_for_branch(repo, branch)
    if data.get("snapshot_sha") != remote_sha or data.get("remote_sha") != remote_sha:
        raise Refusal("preservation receipt does not match canonical archive SHA")
    index_tree=data.get("index_tree")
    if not re.fullmatch(r"[0-9a-f]{40}", index_tree or ""):
        staged_path=data.get("staged_metadata_path")
        try: staged_data=json.loads(Path(staged_path).read_text())
        except (OSError,TypeError,ValueError): staged_data={}
        index_tree=staged_data.get("index_tree")
    if not re.fullmatch(r"[0-9a-f]{40}", index_tree or ""):
        raise Refusal("staged-state metadata lacks original index tree")
    if not run(repo,"cat-file","-e",f"{index_tree}^{{tree}}",check=False).returncode == 0:
        raise Refusal("staged-state metadata index tree unavailable")
    d.parent.mkdir(parents=True,exist_ok=True); run(repo,"worktree","add",str(d),branch)
    run(d,"read-tree",index_tree)
    print(str(d)); return 0

def main(argv):
    if not argv or argv[0] not in {"audit","preserve","retire","restore"}: return die("usage: audit|preserve|retire|restore <archive-branch> <path>")
    if argv[0] == "restore":
        try:
            if len(argv)!=3: raise Refusal("restore requires archive branch and path")
            if any(part==".." for part in Path(argv[2]).parts): raise Refusal("path traversal refused")
            destination=Path(argv[2]).expanduser()
            if path_has_symlink(destination): raise Refusal("restore path contains a symlink")
            repo=root_for(destination.parent)
            return restore(repo,argv[1],argv[2])
        except Refusal as e: return die(str(e))
        except (subprocess.CalledProcessError, OSError) as e: return die(str(e))
    try:
        result=0
        for raw in argv[1:]:
            if ".." in Path(raw).parts: raise Refusal("path traversal refused")
            try:
                repo=root_for(Path(raw).expanduser())
                path,_=target(repo,raw)
                if argv[0]=="audit":
                    entries=inventory(path)
                    counts={kind:sum(e["kind"]==kind for e in entries) for kind in ("staged","unstaged","untracked","generated","secret","ignored-uncertain")}
                    counts["ignored"]=sum(e["kind"]=="ignored" for e in entries)
                    counts["tracked"]=counts["staged"]+counts["unstaged"]
                    print(json.dumps({"path":str(path),"branch":run(path,"branch","--show-current").stdout.strip(),"entries":entries,**counts},sort_keys=True))
                elif argv[0]=="preserve": preserve(repo,path)
                else: retire(repo,path)
            except Refusal as e:
                print(f"archive-worktrees: REFUSED: {raw}: {e}", file=sys.stderr)
                result=2
        return result
    except Refusal as e: return die(str(e))
    except (subprocess.CalledProcessError, OSError) as e: return die(str(e))
if __name__=="__main__": sys.exit(main(sys.argv[1:]))