import gradio as gr
import os
import shutil
import logging
import subprocess
from pathlib import Path
from merge_script import ModelMerger, get_max_vocab_size, download_json_files
import spaces

@spaces.GPU
def merge_models(base_model, model_to_merge, scaling_factor, weight_drop_prob, repo_name, token, commit_message):
    # Define staging and output paths
    staging_path = "/tmp/staging"
    output_path = "/tmp/output"
    os.makedirs(staging_path, exist_ok=True)
    os.makedirs(output_path, exist_ok=True)

    # Initialize ModelMerger and prepare base model
    model_merger = ModelMerger(staging_path, repo_name, token)
    model_merger.prepare_base_model(base_model, os.path.join(staging_path, "base_model"))

    # Merge models and handle progress updates
    for repo_name in [base_model, model_to_merge]:
        model_merger.merge_repo(repo_name, os.path.join(staging_path, "staging_model"), weight_drop_prob, scaling_factor)
        yield 0.25, f"Merged {repo_name}"

    # Finalize merge and handle vocab size
    model_merger.finalize_merge(output_path)
    yield 0.5, "Finalizing merge and handling vocab size..."
    max_vocab_size, repo_with_max_vocab = get_max_vocab_size([base_model, model_to_merge])
    if max_vocab_size > 0:
        download_json_files(repo_with_max_vocab, ['config.json', 'special_tokens_map.json', 'tokenizer.json', 'tokenizer_config.json'], output_path)

    # Upload merged model to Hugging Face Hub
    if repo_name:
        model_merger.upload_model(output_path, repo_name, commit_message)
        yield 0.75, "Uploading merged model to Hugging Face Hub..."
        repo_url = f"https://huggingface.co/{repo_name}"
        yield 1.0, f"Model merged and uploaded successfully! {repo_url}"
    else:
        yield 1.0, "Model merged successfully! No upload performed."

def get_model_type_info(model_name):
    model_types = {
        "base_model": "Base model should be in .safetensors format.",
        "model_to_merge": "Model to merge can be in .safetensors or .bin format."
    }
    return model_types.get(model_name, "No specific info available.")

def validate_model_format(model_name, model_path):
    if model_name == "base_model":
        if not model_path.endswith(".safetensors"):
            return False, "Base model must be in .safetensors format."
    elif model_name == "model_to_merge":
        if not model_path.endswith((".safetensors", ".bin")):
            return False, "Model to merge must be in .safetensors or .bin format."
    return True, None

def merge_and_upload_interface():
    with gr.Blocks(theme="Ytheme/XRainbow", fill_width=True) as demo:
        gr.Markdown("# Model Merger and Uploader")
        gr.Markdown("Combine and upload models with real-time progress updates.")
        gr.Markdown("**Model Compatibility:**")
        gr.Markdown("Combine any two models using a Super Mario merge.")
        gr.Markdown("Works with:")
        gr.Markdown("* Stable Diffusion (1.5, XL/XL Turbo)")
        gr.Markdown("* LLMs (Mistral, Llama, etc)")
        gr.Markdown("* LoRas (must be same size)")
        gr.Markdown("* Any two homologous models")

        with gr.Column():
            token = gr.Textbox(label="HuggingFace Token")
            base_model = gr.Textbox(label="Base Model")
            base_model_info = gr.HTML(get_model_type_info("base_model"))
            model_to_merge = gr.Textbox(label="Model to Merge")
            model_to_merge_info = gr.HTML(get_model_type_info("model_to_merge"))
            repo_name = gr.Textbox(label="New Model Name")
            scaling_factor = gr.Slider(minimum=0, maximum=10, label="Scaling Factor")
            weight_drop_prob = gr.Slider(minimum=0, maximum=1, label="Weight Drop Probability")
            gr.Button("Merge and Upload").click(
                merge_models,
                inputs=[base_model, model_to_merge, scaling_factor, weight_drop_prob, repo_name, token, commit_message],
                outputs=[
                    gr.Textbox(label="Progress/Output"),
                    gr.Textbox(label="Log")
                ],
                pre_process=validate_model_format
            )

    demo.launch()

if __name__ == "__main__":
    merge_and_upload_interface()