#!/usr/bin/env python3
"""Materialize the pinned stable FIPS source plus the reviewed startup patch.

Uses Cargo's existing git checkout offline when available. A fresh workstation
fetches only the pinned public commit. Never edits Cargo's shared source cache.
Run before standalone Cargo commands; Gradle runs it automatically.
"""
import hashlib
import io
import json
import os
from pathlib import Path
import shutil
import subprocess
import tarfile
import tempfile

PIN = 'f4a5b4a8e61e340e6be7551608109a5c399f8353'
URL = 'https://github.com/jmcorgan/fips'
ANDROID = Path(__file__).resolve().parents[1]
PATCH = ANDROID.parent / 'docs/companion-performance/fips-nonblocking-static-discovery.patch'
DEST = ANDROID / 'rust/fips-stable'
STAMP = '.archy-source.json'


def tree_hash(root):
    digest = hashlib.sha256()
    for path in sorted(root.rglob('*')):
        rel = path.relative_to(root).as_posix()
        if rel == STAMP:
            continue
        if path.is_symlink():
            path.resolve().relative_to(root.resolve())
            digest.update(rel.encode() + b'\0link\0' + os.readlink(path).encode())
        elif path.is_file():
            digest.update(rel.encode() + b'\0')
            digest.update(hashlib.sha256(path.read_bytes()).digest())
    return digest.hexdigest()


def main():
    patch_hash = hashlib.sha256(PATCH.read_bytes()).hexdigest()
    expected = {'upstream_commit': PIN, 'patch_sha256': patch_hash}
    if DEST.exists():
        stamp = json.loads((DEST / STAMP).read_text())
        if tree_hash(DEST) != stamp['tree_sha256']:
            raise RuntimeError('Generated FIPS source was edited; preserve edits before regenerating it')
        if all(stamp.get(k) == v for k, v in expected.items()):
            print('Verified pinned FIPS source and startup patch')
            return
    cargo_home = Path(os.environ.get('CARGO_HOME', Path.home() / '.cargo'))
    sources = sorted((cargo_home / 'git/checkouts').glob('fips-*/*'))
    archive = None
    for source in sources:
        check = subprocess.run(['git', '-C', str(source), 'cat-file', '-e', PIN + '^{commit}'],
                               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        if check.returncode == 0:
            archive = subprocess.check_output(['git', '-C', str(source), 'archive', PIN])
            break
    DEST.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix='.prepare-fips-', dir=DEST.parent) as temp:
        temp = Path(temp)
        if archive is None:
            repo = temp / 'upstream.git'
            subprocess.run(['git', 'init', '--bare', '-q', str(repo)], check=True)
            subprocess.run(['git', '--git-dir=' + str(repo), 'fetch', '--depth=1', URL, PIN], check=True)
            actual = subprocess.check_output(['git', '--git-dir=' + str(repo), 'rev-parse', 'FETCH_HEAD'], text=True).strip()
            if actual != PIN:
                raise RuntimeError('Upstream commit mismatch')
            archive = subprocess.check_output(['git', '--git-dir=' + str(repo), 'archive', PIN])
        stage = temp / 'source'
        stage.mkdir()
        with tarfile.open(fileobj=io.BytesIO(archive)) as tar:
            tar.extractall(stage, filter='data')
        # Isolate git-apply from the enclosing application's repository prefix.
        subprocess.run(['git', 'init', '-q', str(stage)], check=True)
        subprocess.run(['git', 'apply', '--check', str(PATCH)], cwd=stage, check=True)
        subprocess.run(['git', 'apply', str(PATCH)], cwd=stage, check=True)
        shutil.rmtree(stage / '.git')
        expected['tree_sha256'] = tree_hash(stage)
        (stage / STAMP).write_text(json.dumps(expected, indent=2) + '\n')
        if DEST.exists():
            shutil.rmtree(DEST)  # verified generated tree only, never user source
        stage.rename(DEST)
    print('Prepared stable FIPS ' + PIN + ' with reviewed nonblocking discovery patch')


if __name__ == '__main__':
    main()