#!/usr/bin/env python3
# Compiled-shader cache keyed by the Metal source (dmsh_build.py): a library is compiled once per (MSL text, compile convention, compiler), so a chain change
# that leaves a shader's MSL as it was (a translator fix, a new rule for other shaders) costs that shader no Metal compile. The Metal compiler is deterministic:
# the same MSL with the same flags gives a byte-identical metallib (checked for macOS and iOS), which is what makes reuse safe. Its input is the text alone: dmsh_pch.compile passes the source on
# standard input, as a file's name reaches the library of a source with function constants.
# Entries of the compile cache's one file (shader_cache.py): metallib/<target>/<digest>.metallib, digest = sha256 over FORMAT, the target's compile convention (the record written into json chain_metal),
# the compiler's identity (metal / metallib --version without the install path, the SDK version and build) and the MSL bytes.
# Filling from earlier builds: every imported library must record the current compiler identity and flags in its reflection. Legacy libraries without that identity are not imported.
# A fresh compile of two shaders per stage also checks their bytes. Each folder is read once per compiler identity (filled.json remembers its number of libraries).
# Not part of the chain version (dmsh_chain.py): the cache never changes an output, it only skips recompiling one. Change FORMAT when what it stores changes.
import os, re, sys, json, shutil, ctypes, hashlib, threading, subprocess, collections
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', 'tools'))
import romfs, shader_lib   # the libraries of an earlier build may be one file (shader_lib.py), read in place
STAMPS = 'metallib/filled.json'   # the earlier builds' folders already read

FORMAT = b'dmsh_mtlcache 2'   # earlier prefill could misclassify an unsampled library from another compiler
OUTPUT = re.compile(r'[0-9a-f]{16}_b\.metal')   # a shader's MSL in dmsh_build.py's outputs: <key>_b.metal
_clonefile = getattr(ctypes.CDLL(None, use_errno=True), 'clonefile', None)   # APFS copy-on-write copy (no data written); elsewhere a plain copy


def clone(src, dst):
    tmp = f'{dst}.tmp-{os.getpid()}-{threading.get_ident()}'
    if _clonefile is None or _clonefile(os.fsencode(src), os.fsencode(tmp), 0) != 0: shutil.copyfile(src, tmp)
    os.replace(tmp, dst)


def compiler_id(sdk):
    def run(*a): return subprocess.run(['xcrun', '-sdk', sdk, *a], capture_output=True, text=True).stdout
    tools = ''.join(run(t, '--version') for t in ('metal', 'metallib'))
    return re.sub(r'(?m)^InstalledDir:.*\n?', '', tools) + run('--show-sdk-version') + run('--show-sdk-build-version')   # the install path changes with each toolchain mount


class Cache:
    def __init__(self, store, targets, work):
        """store: shader_cache.Store; targets: {name: (sdk, record, compile)}, compile(msl, air, lib) as ios_metal.compile; work: a folder for temporary .air files"""
        self.store, self.targets, self.work = store, targets, work
        self.ident = {t: hashlib.sha256(FORMAT + b'\0' + rec.encode() + b'\0' + compiler_id(sdk).encode()).digest() for t, (sdk, rec, _) in targets.items()}
        self.stats, self.lock = collections.Counter(), threading.Lock()

    def path(self, target, src):
        d = hashlib.sha256(self.ident[target] + b'\0' + src).hexdigest()
        return f'metallib/{target}/{d}.metallib'

    def with_compiler(self, target, refl):
        return dict(refl, chain_metal=self.targets[target][1], chain_compiler=self.ident[target].hex())

    def count(self, what, n=1):
        with self.lock: self.stats[what] += n

    def _compile(self, target, msl, lib):
        air = os.path.join(self.work, f'mtlcache-{target}-{os.getpid()}-{threading.get_ident()}.air')
        try: self.targets[target][2](msl, air, lib)
        finally:
            if os.path.exists(air): os.unlink(air)

    def compile(self, target, msl, lib):
        """msl → lib for target: a clone of the library that the same source compiled to before, or a compile (then kept). Raises CalledProcessError as the compile does"""
        p = self.path(target, open(msl, 'rb').read())
        if self.store.to_file(p, lib): self.count((target, 'reused')); return
        tmp = f'{lib}.tmp-{os.getpid()}-{threading.get_ident()}'
        self._compile(target, msl, tmp)
        os.replace(tmp, lib)
        self.store.put_file(p, lib); self.count((target, 'compiled'))

    def compile_split(self, target, msl, libs, split):
        """msl → the libraries libs cut from its one compile, split(the library's bytes) → their bytes, in order: clones of what the same source gave before, or a compile (then the pieces
        kept, not the whole library, so an output stays a clone as compile's are). The pieces are cached by their place (<digest>.metallib, <digest>_x.metallib, <digest>_x2.metallib):
        one source always splits the same way. Raises CalledProcessError as the compile does, and what split raises"""
        p = self.path(target, open(msl, 'rb').read())
        paths = [p] + [f'{p[:-len(".metallib")]}_x{i if i > 1 else ""}.metallib' for i in range(1, len(libs))]
        if all(self.store.has(x) for x in paths):
            for src, dst in zip(paths, libs): self.store.to_file(src, dst)
            self.count((target, 'reused')); return
        whole = f'{libs[0]}.tmp-{os.getpid()}-{threading.get_ident()}'
        self._compile(target, msl, whole)
        try:
            with open(whole, 'rb') as f: pieces = split(f.read())
        finally: os.unlink(whole)
        if len(pieces) != len(libs): raise ValueError(f'the split gave {len(pieces)} libraries, not {len(libs)}')
        for lib, data in zip(libs, pieces):
            with open(whole, 'wb') as f: f.write(data)
            os.replace(whole, lib)
        with self.store.batch():
            for path, lib in zip(paths, libs): self.store.put_file(path, lib)
        self.count((target, 'compiled'))

    def _check(self, target, d, stems, mac_dir):
        """None if folder d was built with the current convention and compiler, else why not"""
        rec, picks = self.targets[target][1], collections.defaultdict(list)
        for stem in stems:
            try: j = json.loads(romfs.open(os.path.join(d, stem + '.json'), 'rb').read())
            except (OSError, ValueError): return f'no readable {stem}.json'
            if j.get('chain_metal') != rec: return f'built with "{j.get("chain_metal")}"'
            if j.get('chain_compiler') != self.ident[target].hex(): return f'{stem} has a different or unrecorded Metal compiler'
            if len(picks[j.get('stage')]) < 2: picks[j.get('stage')].append(stem)
        for stem in (s for v in picks.values() for s in v):
            tmp = os.path.join(self.work, f'mtlcache-check-{target}.metallib')
            msl = os.path.join(mac_dir, stem + '.metal')
            inside = mac_dir.endswith('.img')   # the compiler is then given a file of its own, which goes with the check
            if inside:
                with romfs.open(msl, 'rb') as f, open(os.path.join(self.work, f'mtlcache-check-{stem}.metal'), 'wb') as o: shutil.copyfileobj(f, o)
                msl = os.path.join(self.work, f'mtlcache-check-{stem}.metal')
            try: self._compile(target, msl, tmp)
            except subprocess.CalledProcessError: return f'{stem}.metal does not compile now'
            finally:
                if inside: os.unlink(msl)
            same = open(tmp, 'rb').read() == romfs.open(os.path.join(d, stem + '.metallib'), 'rb').read(); os.unlink(tmp)
            if not same: return f'a fresh compile of {stem}.metal differs from its library'
        return None

    def fill(self, mac_dir, ios_dir=None, targets=None):
        """Takes the libraries of an earlier build: mac_dir = its Mac b-<chain>/ folder, ios_dir = the iOS folder of the same chain (or None). Returns the lines to print"""
        mac_dir, ios_dir = str(shader_lib.view(mac_dir)), ios_dir and str(shader_lib.view(ios_dir))
        if not romfs.os.path.isdir(mac_dir): return []   # a stopped pack's .part beside the libraries
        try: stamps = json.loads(self.store.get(STAMPS) or b'{}')
        except ValueError: stamps = {}
        names = set(romfs.os.listdir(mac_dir))   # a library with a side or readable-tile library was cut from another source than its .metal (compile_split, dmsh_side.py), and a compile of that source is never asked for
        stems = sorted(n[:-6] for n in names if OUTPUT.fullmatch(n) and n[:-6] + 'x.metallib' not in names and n[:-6] + 'f.metallib' not in names)   # not an AppleDouble ._<name> (exFAT / FAT drives)
        todo, notes = {}, []
        for target, d in (('mac', mac_dir), ('ios', ios_dir)):
            if targets is not None and target not in targets: continue   # checking an earlier build compiles samples too
            if target not in self.targets or not d or not romfs.os.path.isdir(d): continue
            libs = names if d == mac_dir else set(romfs.os.listdir(d))
            have = [s for s in stems if s + '.metallib' in libs]; key = self.ident[target].hex() + ':' + os.path.realpath(d)
            if not have or stamps.get(key) == len(have): continue
            why = self._check(target, d, have, mac_dir)
            if why: notes.append(f'Metal cache: not filled from {d}: {why}')
            else: todo[target] = (d, set(have))
            stamps[key] = len(have)
        added = collections.Counter()
        for stem in stems:
            ts = [t for t, (d, have) in todo.items() if stem in have]
            if not ts: continue
            src = romfs.open(os.path.join(mac_dir, stem + '.metal'), 'rb').read()
            for t in ts:
                p = self.path(t, src)
                if self.store.has(p): continue
                self.store.put(p, romfs.open(os.path.join(todo[t][0], stem + '.metallib'), 'rb').read()); added[t] += 1
        for t, (d, have) in todo.items():
            self.count((t, 'filled'), added[t]); notes.append(f'Metal cache: {added[t]} {t} libraries added from {d} ({len(have)} there)')
        self.store.put(STAMPS, json.dumps(stamps, indent=1).encode())
        return notes

    def summary(self):
        s = self.stats
        return 'Metal cache: ' + '; '.join(f'{t} {s[t, "reused"]} reused, {s[t, "compiled"]} compiled' + (f', {s[t, "filled"]} filled from earlier builds' if s[t, 'filled'] else '')
                                          for t in self.targets)