#!/usr/bin/env python3
"""
Bundle a Strategy-1 KV-cache TFLite + SentencePiece tokenizer into a
.litertlm file compatible with Google AI Edge / LiteRT-LM runtime.

Embeds:
  - LlmMetadata proto: Gemma3 model type, 2K max tokens, TranslateGemma
    Jinja chat template, BOS/EOS/end_of_turn stop tokens
  - TFLite model  (model_type=prefill_decode)
  - SentencePiece tokenizer

Usage:
  python bundle_litertlm.py \
    --tflite /path/to/model.tflite \
    --tokenizer /path/to/tokenizer.model \
    --output /path/to/output.litertlm \
    [--max-tokens 2048]
"""

import argparse
import sys
import tempfile
from pathlib import Path

# Make litert_lm package importable from /tmp/litert-lm-pkg
sys.path.insert(0, "/tmp/litert-lm-pkg")

from litert_lm_builder import litertlm_builder
from litert_lm.runtime.proto import (
    llm_metadata_pb2,
    llm_model_type_pb2,
    token_pb2,
)


# Simple Jinja template compatible with LiteRT-LM runtime (no .get(), no complex tests).
# Handles plain text input from Google AI Edge Gallery.
# Uses the exact prompt format TranslateGemma was trained with (en→es default).
# Users who need other language pairs should prefix their message with the pair,
# e.g. "Translate English to French:\n\nHello"
TRANSLATE_GEMMA_JINJA_TEMPLATE = \
"{{ bos_token }}" \
"{% for message in messages %}" \
"{% if message['role'] == 'user' %}" \
"<start_of_turn>user\n" \
"You are a professional translator. " \
"Produce only the translation of the following text, without any additional explanations or commentary:\n\n\n" \
"{{ message['content'] | trim }}" \
"<end_of_turn>\n" \
"{% elif message['role'] == 'assistant' %}" \
"<start_of_turn>model\n" \
"{{ message['content'] | trim }}" \
"<end_of_turn>\n" \
"{% endif %}" \
"{% endfor %}" \
"{% if add_generation_prompt %}" \
"<start_of_turn>model\n" \
"{% endif %}"


def build_llm_metadata_proto(max_tokens: int) -> bytes:
    meta = llm_metadata_pb2.LlmMetadata()
    meta.max_num_tokens = max_tokens

    # Model type: Gemma3 (text-only variant — no vision config needed for TranslateGemma text mode)
    meta.llm_model_type.gemma3.CopyFrom(llm_model_type_pb2.Gemma3())

    # Start token: BOS = token id 2
    meta.start_token.token_ids.ids.append(2)

    # Stop tokens: EOS (id=1) and end_of_turn (id=106)
    eos = meta.stop_tokens.add()
    eos.token_ids.ids.append(1)

    eot = meta.stop_tokens.add()
    eot.token_ids.ids.append(106)

    # Embed the Jinja template
    meta.jinja_prompt_template = TRANSLATE_GEMMA_JINJA_TEMPLATE

    return meta.SerializeToString()


def main():
    ap = argparse.ArgumentParser(description="Bundle TFLite + tokenizer into .litertlm")
    ap.add_argument("--tflite", required=True)
    ap.add_argument("--tokenizer", required=True, help="SentencePiece .model file")
    ap.add_argument("--output", required=True)
    ap.add_argument("--max-tokens", type=int, default=2048)
    ap.add_argument("--quant", default="int8", help="Quantization label for metadata")
    args = ap.parse_args()

    tflite_path = Path(args.tflite)
    tokenizer_path = Path(args.tokenizer)
    output_path = Path(args.output)

    if not tflite_path.exists():
        print(f"[x] TFLite not found: {tflite_path}", file=sys.stderr)
        sys.exit(1)
    if not tokenizer_path.exists():
        print(f"[x] Tokenizer not found: {tokenizer_path}", file=sys.stderr)
        sys.exit(1)

    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Write LlmMetadata to temp file
    meta_bytes = build_llm_metadata_proto(args.max_tokens)
    with tempfile.NamedTemporaryFile(suffix=".pb", delete=False) as f:
        meta_file = Path(f.name)
        f.write(meta_bytes)

    print(f"[+] Building .litertlm: {output_path.name}")
    print(f"    TFLite: {tflite_path} ({tflite_path.stat().st_size / 1e9:.2f} GB)")
    print(f"    Tokenizer: {tokenizer_path}")
    print(f"    Max tokens: {args.max_tokens}")

    Metadata = litertlm_builder.Metadata
    DType = litertlm_builder.DType

    builder = litertlm_builder.LitertLmFileBuilder()
    builder.add_system_metadata(Metadata(key="model_name", value=f"TranslateGemma-4B-IT-{args.quant}", dtype=DType.STRING))
    builder.add_system_metadata(Metadata(key="authors", value="google", dtype=DType.STRING))
    builder.add_system_metadata(Metadata(key="quantization", value=args.quant, dtype=DType.STRING))

    builder.add_tflite_model(
        str(tflite_path),
        model_type=litertlm_builder.TfLiteModelType.PREFILL_DECODE,
    )
    builder.add_sentencepiece_tokenizer(str(tokenizer_path))
    builder.add_llm_metadata(str(meta_file))

    with open(output_path, "wb") as f:
        builder.build(f)
    meta_file.unlink(missing_ok=True)

    size = output_path.stat().st_size
    print(f"[+] Written: {output_path} ({size / 1e9:.2f} GB)")


if __name__ == "__main__":
    main()