# Copyright 2026 Spawn Lau
# SPDX-License-Identifier: Apache-2.0
"""Refine the proven 630-entry candidate to the native 70-token grid bound."""

import argparse
import hashlib
import json
from pathlib import Path

from ai_edge_litert import schema_py_generated as schema
from litert_lm_builder import LitertLmFileBuilder

from build_size_candidates import crop_positions, pack_model


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--candidate-dir", type=Path, required=True)
    args = parser.parse_args()
    unpacked = args.candidate_dir / "unpacked"
    report_path = args.candidate_dir / "build-report.json"
    report = json.loads(report_path.read_text(encoding="utf-8"))
    assert not any(x["name"] == "position210" for x in report["candidates"])
    model_path, = unpacked.glob("*TFLiteModel*tf_lite_vision_encoder*.tflite")
    weight_path, = unpacked.glob("*TFLiteWeights*tf_lite_vision_encoder*.weight")
    builder = LitertLmFileBuilder.from_toml_file(str(unpacked / "model.toml"))
    model = schema.ModelT.InitFromObj(schema.Model.GetRootAsModel(model_path.read_bytes(), 0))
    refinement = crop_positions(model, weight_path, 210)
    refinement["native_bound_argument"] = {
        "patch_size": 16, "pooling_kernel_size": 3, "max_subpatches": 630,
        "preprocessed_dimensions_are_multiples_of_pixels": 48,
        "minimum_patch_grid_side": 3,
        "maximum_patch_grid_side": 210,
        "valid_position_coordinate_range": [0, 209],
        "source": "https://github.com/google-ai-edge/LiteRT-LM/blob/v0.18.0/support/preprocessor/image_preprocessor_utils.cc",
        "scope": "Current app's native image preprocessing with vision_tokens_per_image=70 only.",
    }
    model_path.write_bytes(pack_model(model))
    output_path = args.candidate_dir / "eg2-gallery-position210.litertlm"
    with output_path.open("wb") as destination:
        builder.build(destination)
    report["position210_refinement"] = refinement
    candidate = {"name": "position210", "path": str(output_path), "bytes": output_path.stat().st_size,
                 "saved_bytes": report["source_bytes"] - output_path.stat().st_size,
                 "sha256": hashlib.file_digest(output_path.open("rb"), "sha256").hexdigest()}
    report["candidates"].append(candidate)
    report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps(candidate, indent=2))


if __name__ == "__main__":
    main()