import os
import logging
from litert_torch import converter
from mediapipe.tasks.python.genai import bundler

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

MODEL_PATH = "artifacts/deploy/ankahi-gemma4-e4b-int8"
TFLITE_PATH = "artifacts/litertlm/ankahi_q8.tflite"
TASK_PATH = "artifacts/litertlm/ankahi_q8.task"

def main():
    os.makedirs("artifacts/litertlm", exist_ok=True)
    
    log.info(f"Converting {MODEL_PATH} to LiteRT (.tflite)...")
    try:
        converter.convert(
            model_path=MODEL_PATH,
            output_path=TFLITE_PATH,
            quantize="INT8",
            prefill_seq_len=512,  # Reduced for memory/latency
            kv_cache_max_len=1024,
            backend="cpu"
        )
        log.info(f"Successfully created TFLite model at {TFLITE_PATH}")
    except Exception as e:
        log.error(f"LiteRT conversion failed: {e}")
        return

    log.info(f"Bundling into MediaPipe Task (.task)...")
    try:
        # Check for tokenizer.model or tokenizer.json
        tokenizer_model = os.path.join(MODEL_PATH, "tokenizer.model")
        if not os.path.exists(tokenizer_model):
            # Fallback or check if tokenizer.json is used
            log.warning("tokenizer.model not found, searching for alternatives...")
            tokenizer_model = os.path.join(MODEL_PATH, "tokenizer.json")
            
        config = bundler.BundleConfig(
            tflite_model=TFLITE_PATH,
            tokenizer_model=tokenizer_model,
            start_token="<bos>",
            stop_tokens=["<eos>", "<end_of_turn>"],
            output_filename=TASK_PATH,
            enable_streaming=True
        )
        bundler.create_bundle(config)
        log.info(f"Successfully created MediaPipe task at {TASK_PATH}")
    except Exception as e:
        log.error(f"Task bundling failed: {e}")

if __name__ == "__main__":
    main()