"""Round 2 step 3: fp32 export at any contract shape (export_s128.py with the shape as an argument).

Usage: export_shape.py --seq 256   (artifact exports/gliner25_decide_s{N}_fp32.tflite; an existing
artifact is never overwritten). Sample inputs = the first oracle fixture that fits s{N}.
Writes results/opscan_s{N}.json (op histogram, ranks, BATCH_MATMUL shapes, forbidden check) and
results/export_s{N}.json (bytes, sha256, export seconds, op/forbidden summary). The forbidden check
is round 1's `forbidden_check`, imported unchanged.
"""
import argparse
import hashlib
import json
import time
import traceback

import torch

from common import ROOT, dump_json, load_model
from decide_graph import DecideClassifier, Shape, fixed_inputs
from export_s128 import forbidden_check


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--seq", type=int, required=True)
    args = parser.parse_args()
    seq = args.seq
    artifact = ROOT / "exports" / f"gliner25_decide_s{seq}_fp32.tflite"
    assert not artifact.exists(), "Never overwrite an exported artifact"
    for relative in (f"results/export_s{seq}.json", f"results/opscan_s{seq}.json"):
        assert not (ROOT / relative).exists(), f"Refusing to overwrite {relative}"
    shape = Shape(seq)
    oracle = json.loads((ROOT / "fixtures/oracle_fp32.json").read_text())
    sample = next(f for f in oracle["fixtures"] if f["fits"][f"s{seq}"])
    model = load_model()
    graph = DecideClassifier(model, shape).eval().requires_grad_(False)
    captured = torch.load(ROOT / f"fixtures/captured/{sample['id']}.pt", weights_only=False)
    inputs = fixed_inputs(model, captured, shape)
    entry = {"artifact": str(artifact.relative_to(ROOT)), "sample_fixture": sample["id"],
             "contract": shape.contract(), "status": "RUNNING",
             "converter": "litert-torch (venv lock) + converter_layout.register_rank4_matmul"}
    started = time.perf_counter()
    try:
        import litert_torch
        from converter_layout import register_rank4_matmul
        from litert_cpu import scan
        register_rank4_matmul()
        litert_torch.convert(graph, inputs).export(str(artifact))
        elapsed = time.perf_counter() - started
        digest = hashlib.sha256(artifact.read_bytes()).hexdigest()
        opscan = scan(artifact)
        opscan.update(forbidden_check(artifact))
        dump_json(f"results/opscan_s{seq}.json", opscan)
        shapes = {}
        for item in opscan["batch_matmuls"]:
            key = " x ".join(str(s) for s in item["input_shapes"])
            shapes[key] = shapes.get(key, 0) + 1
        entry.update(status="PASS", bytes=artifact.stat().st_size, sha256=digest, export_seconds=elapsed,
                     operator_count=opscan["operator_count"], op_histogram=opscan["op_histogram"],
                     max_tensor_rank=opscan["max_tensor_rank"], tensor_dtype_histogram=opscan["tensor_dtype_histogram"],
                     all_batch_matmul_inputs_rank4=opscan["all_batch_matmul_inputs_rank4"],
                     batch_matmul_input_shapes=shapes, forbidden_checked=opscan["checked"],
                     forbidden_count=opscan["forbidden_count"], forbidden_found=opscan["forbidden_found"],
                     signatures=opscan["signatures"], opscan=f"results/opscan_s{seq}.json")
        dump_json(f"results/export_s{seq}.json", entry)
        print(json.dumps({k: v for k, v in entry.items() if k not in ("contract", "signatures")}, indent=1), flush=True)
    except BaseException:
        log = ROOT / f"logs/export_{artifact.stem}.traceback.txt"
        log.write_text(traceback.format_exc())
        entry.update(status="FAIL", export_seconds=time.perf_counter() - started, traceback=str(log.relative_to(ROOT)))
        dump_json(f"results/export_attempt_{artifact.stem}.json", entry)
        raise


if __name__ == "__main__":
    main()