#!/usr/bin/env python3
"""Ensambla un weights.rs desde varios certificados del banco de referencia (Laurha, la Pi 5).

    python3 scripts/bench/ensamblar-pesos.py <weights.rs> <cert-antiguo.rs> ... <cert-vigente.rs>
    python3 scripts/bench/ensamblar-pesos.py --comprobar <weights.rs> <certs...>   # 1 si no coincide

Por qué existe: «El máximo de 5 de un día no es el de otro» (Laurha, 29-oct). Entre dos
certificados del mismo código, el base weight de un mismo extrinsic se separa hasta un 22 %. Si se
cobra sólo el último certificado, se cobra de menos lo que otro día midió más alto sobre el MISMO
código. Si se mezclan certificados de código distinto, se cobra lo que no corre.

La regla, función a función (en los dos impl, `SubstrateWeight<T>` y `()`):
- El ÚLTIMO certificado de la lista es el del código que corre. Su bloque es la referencia: sus
  lecturas, escrituras, proof y lista de storage son las que se cobran.
- Entran todos los certificados cuyo bloque tiene el MISMO esqueleto que la referencia: el cuerpo
  sin números ni comentarios numéricos. Un certificado de código anterior que lea otra cosa no entra.
- Cada `Weight::from_parts(a, b)` cobra el máximo de los que entran, coeficiente a coeficiente.
- Se niega si el esqueleto de la referencia no es el que ya tiene el weights.rs (el código cambió y
  no hay certificado de él) o si una función del weights.rs no está en la referencia.

Sólo toca los bloques de las funciones. La cabecera, el trait y la guarda de proof_size se quedan
como están: la cabecera la escribe quien ensambla, y debe nombrar los certificados y sus shas.
"""
import hashlib, os, re, sys

IMPLS = ("impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {", "impl WeightInfo for () {")
FN = re.compile(r"((?:\t//[^\n]*\n)*)\tfn (\w+)\(\) -> Weight \{\n.*?\n\t\}\n", re.S)
PARTES = re.compile(r"from_parts\(([\d_]+), ([\d_]+)\)")


def impl(texto, cabeza):
    a = texto.index(cabeza)
    b = texto.index("\n}\n", a) + 1
    return a, b


def bloques(texto, cabeza):
    a, b = impl(texto, cabeza)
    return {m.group(2): m for m in FN.finditer(texto, a, b)}


def esqueleto(bloque):
    cuerpo = bloque.split("\tfn ", 1)[1]
    cuerpo = re.sub(r"\t\t// [^\n]*\n", "", cuerpo)
    return PARTES.sub("from_parts(N, N)", cuerpo)


def numeros(bloque):
    return [(int(a.replace("_", "")), int(b.replace("_", ""))) for a, b in PARTES.findall(bloque)]


def ensamblar(destino, certs):
    texto = open(destino).read()
    leidos = [(os.path.basename(c).rsplit(".", 1)[0], open(c).read()) for c in certs]
    for cabeza in IMPLS:
        actuales = bloques(texto, cabeza)
        ref = bloques(leidos[-1][1], cabeza)
        otros = [(et, bloques(t, cabeza)) for et, t in leidos]
        nuevos = {}
        for nombre, m in actuales.items():
            if nombre not in ref:
                sys.exit(f"{nombre}: no está en el certificado vigente ({leidos[-1][0]})")
            r = ref[nombre].group(0)
            if esqueleto(r) != esqueleto(m.group(0)):
                sys.exit(f"{nombre}: el certificado vigente no mide el código del weights.rs")
            entran = [(et, b[nombre].group(0)) for et, b in otros
                      if nombre in b and esqueleto(b[nombre].group(0)) == esqueleto(r)]
            maximos = [tuple(max(v) for v in zip(*col)) for col in zip(*[numeros(b) for _, b in entran])]
            it = iter(maximos)
            bloque = PARTES.sub(lambda _m: "from_parts({:_}, {:_})".format(*next(it)), r)
            etiquetas = ", ".join(et for et, _ in entran)
            bloque = re.sub(r"\t\t// Tiempo minimo de ejecucion:[^\n]*\n",
                            f"\t\t// Máximo por coeficiente de: {etiquetas}.\n", bloque)
            nuevos[nombre] = bloque
        a, b = impl(texto, cabeza)
        seccion = FN.sub(lambda mm: nuevos[mm.group(2)], texto[a:b])
        texto = texto[:a] + seccion + texto[b:]
    return texto


if __name__ == "__main__":
    args = sys.argv[1:]
    comprobar = bool(args) and args[0] == "--comprobar"
    if comprobar:
        args = args[1:]
    if len(args) < 2:
        sys.exit(__doc__)
    destino, certs = args[0], args[1:]
    nuevo = ensamblar(destino, certs)
    for c in certs:
        print(f"  {hashlib.sha256(open(c, 'rb').read()).hexdigest()[:16]}  {c}")
    if comprobar:
        igual = nuevo == open(destino).read()
        print("weights.rs coincide con el ensamblado" if igual else "weights.rs NO coincide con el ensamblado")
        sys.exit(0 if igual else 1)
    open(destino, "w").write(nuevo)
    print(f"ensamblado {destino} desde {len(certs)} certificados")