# app.py
import gradio as gr
from optimum.exporters.coreml import main_export
import os
import shutil
import zipfile

# --- This is the model we will download from the "Supplier" ---
MODEL_TO_CONVERT = "sentence-transformers/all-MiniLM-L6-v2"
MODEL_NAME = "MiniLM-L6-v2-CoreML" # Our new name for the converted model

# --- This is the main conversion function ---
def run_conversion():
    print("--- Starting Conversion ---")
    
    try:
        # 1. Export the model
        # This is the main command that uses 'optimum' to convert
        # the model from Hugging Face to the Core ML format.
        main_export(
            model_name_or_path=MODEL_TO_CONVERT,
            output=MODEL_NAME,
            task="feature-extraction", # This is the correct task for embedding
            do_quantize=False # We will not quantize it for now
        )
        
        print(f"--- Model exported to folder: {MODEL_NAME} ---")

        # 2. Zip the converted model folder
        zip_path = f"{MODEL_NAME}.zip"
        shutil.make_archive(base_name=MODEL_NAME,
                            format='zip',
                            root_dir='.',
                            base_dir=MODEL_NAME)
        
        print(f"--- Model zipped to: {zip_path} ---")
        
        # 3. Return the path to the zip file so the user can download it
        return f"Success! Converted model is ready. Download it here: <a href='/file={zip_path}' target='_blank'>{zip_path}</a>"

    except Exception as e:
        print(f"--- ERROR: {e} ---")
        return f"Error during conversion: {e}"

# --- This creates the simple web UI ---
with gr.Blocks() as demo:
    gr.Markdown("# CortexApp Model Converter (MiniLM-L6-v2)")
    gr.Markdown("Click the button to convert the `sentence-transformers/all-MiniLM-L6-v2` model to Core ML format and zip it for download.")
    
    start_button = gr.Button("Start Conversion")
    output_text = gr.HTML("Conversion has not started.")
    
    start_button.click(
        fn=run_conversion,
        inputs=[],
        outputs=[output_text]
    )

demo.launch()