# Copyright 2026 Spawn Lau
# SPDX-License-Identifier: Apache-2.0
# Release change: optional fixed container metadata for byte-exact rebuilding.
"""Remove unreferenced buffer records and deduplicate identical inline data."""

import argparse
import hashlib
import json
from pathlib import Path

import flatbuffers
import numpy as np
from ai_edge_litert import schema_py_generated as schema
from litert_lm_builder import DType, LitertLmFileBuilder, Metadata
from litert_lm_builder import litertlm_builder as builder_impl


def published_system_metadata(entries):
    # LiteRT-LM builder 0.18.0 normally generates a new UUID and timestamp on
    # every pack. Pin these non-computational fields to the tested release.
    return [entry for entry in entries if entry.key not in {"uuid", "creation_timestamp"}] + [
        Metadata(key="uuid", value="d8074fb5-3ec2-4550-9558-2c103b3ab73e", dtype=DType.STRING),
        Metadata(key="creation_timestamp", value="2026-10-20T12:27:29.752155+00:00", dtype=DType.STRING),
    ]


class SharedStringBuilder(flatbuffers.Builder):
    def CreateString(self, value, encoding="utf-8", errors="strict"):
        encoded = value.encode(encoding, errors) if isinstance(value, str) else bytes(value)
        if self.sharedStrings is None:
            self.sharedStrings = {}
        if encoded not in self.sharedStrings:
            self.sharedStrings[encoded] = super().CreateString(encoded)
        return self.sharedStrings[encoded]


def pack(model):
    builder = SharedStringBuilder(1024)
    builder.Finish(model.Pack(builder), file_identifier=b"TFL3")
    return bytes(builder.Output())


def compact(path):
    raw = path.read_bytes()
    model = schema.ModelT.InitFromObj(schema.Model.GetRootAsModel(raw, 0))
    before_buffers = len(model.buffers)
    before_external = len(model.externalBuffers or [])
    used = {0}
    used_external = set()
    for graph in model.subgraphs:
        for tensor in graph.tensors:
            assert not tensor.isVariable
            used.add(int(tensor.buffer))
            if tensor.externalBuffer:
                used_external.add(int(tensor.externalBuffer))
    for metadata in model.metadata or []:
        used.add(int(metadata.buffer))
    used.update(int(i) for i in model.metadataBuffer or [])
    mapping, keys, buffers = {}, {}, []
    for old in sorted(used):
        buffer = model.buffers[old]
        key = (bytes(buffer.data) if buffer.data is not None else b"", int(buffer.offset), int(buffer.size))
        if key not in keys:
            keys[key] = len(buffers)
            buffers.append(buffer)
        mapping[old] = keys[key]
    for graph in model.subgraphs:
        for tensor in graph.tensors:
            old_buffer = model.buffers[int(tensor.buffer)]
            new_buffer = buffers[mapping[int(tensor.buffer)]]
            assert bytes(old_buffer.data or b"") == bytes(new_buffer.data or b"") if not isinstance(old_buffer.data, np.ndarray) else np.array_equal(old_buffer.data, new_buffer.data)
            tensor.buffer = mapping[int(tensor.buffer)]
    for metadata in model.metadata or []:
        metadata.buffer = mapping[int(metadata.buffer)]
    if model.metadataBuffer is not None:
        model.metadataBuffer = np.asarray([mapping[int(i)] for i in model.metadataBuffer], dtype=np.int32)
    model.buffers = buffers
    model.externalBuffers = [e for e in model.externalBuffers or [] if int(e.id) in used_external]
    result = pack(model)
    path.write_bytes(result)
    return {"file": path.name, "before_graph_bytes": len(raw), "after_graph_bytes": len(result),
            "before_buffer_records": before_buffers, "after_buffer_records": len(buffers),
            "before_external_buffer_records": before_external,
            "after_external_buffer_records": len(model.externalBuffers),
            "tensor_names_retained": True, "referenced_constant_values_retained": True}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--candidate-dir", type=Path, required=True)
    parser.add_argument("--reproducible", action="store_true",
                        help="Use the published container UUID and timestamp.")
    args = parser.parse_args()
    if args.reproducible:
        builder_impl.populate_system_metadata = published_system_metadata
    report_path = args.candidate_dir / "build-report.json"
    report = json.loads(report_path.read_text(encoding="utf-8"))
    assert not any(c["name"] == "compact_final" for c in report["candidates"])
    unpacked = args.candidate_dir / "unpacked"
    builder = LitertLmFileBuilder.from_toml_file(str(unpacked / "model.toml"))
    details = [compact(path) for path in sorted(unpacked.glob("*.tflite"))]
    output = args.candidate_dir / "eg2-gallery-compact.litertlm"
    with output.open("wb") as target:
        builder.build(target)
    candidate = {"name": "compact_final", "path": str(output), "bytes": output.stat().st_size,
                 "saved_bytes": report["source_bytes"] - output.stat().st_size,
                 "sha256": hashlib.file_digest(output.open("rb"), "sha256").hexdigest()}
    report["metadata_compaction"] = details
    report["candidates"].append(candidate)
    report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps({"candidate": candidate, "details": details}, indent=2))


if __name__ == "__main__":
    main()