import subprocess
import gradio as gr
from huggingface_hub import HfApi


def install_git_repo_editable(repo_url, target_dir):
    """
    Clones a git repository and installs it in editable mode.
    """
    try:
        # 1. Clone the repository
        print(f"Cloning {repo_url} into {target_dir}...")
        subprocess.run(["git", "clone", repo_url, target_dir], check=True)
        
        # 2. Install in editable mode
        print(f"Installing {target_dir} in editable mode...")
        # Using 'python -m pip' is generally safer than calling 'pip' directly
        subprocess.run(["python", "-m", "pip", "install", "-e", target_dir], check=True)
        
        print("Installation successful.")
    except subprocess.CalledProcessError as e:
        print(f"An error occurred: {e}")

def run_merge_and_upload(repo_id, config_path, output_path, progress=gr.Progress()):
    logs = ""

    # 1. Start Merge
    progress(0, desc="Starting merge...")
    logs += f"Running: mergekit-moe-qwen2 {config_path} {output_path}\n"

    command = ["python", "qwen2_moe.py", config_path, output_path]
    try:
        # We use Popen to capture logs if needed, but for simplicity:
        result = subprocess.run(command, capture_output=True, text=True, check=True)
        logs += "Merge successful!\n"
    except subprocess.CalledProcessError as e:
        return f"Error during merge:\n{e.stderr}", "Failed"

    # 2. Start Upload
    progress(0.5, desc="Uploading to Hugging Face...")
    logs += f"Uploading folder {output_path} to {repo_id}...\n"

    try:
        api = HfApi()
        api.upload_folder(
            folder_path=output_path,
            repo_id=repo_id,
            repo_type="model"
        )
        logs += f"Done! Model available at hf.co/{repo_id}"
        progress(1.0, desc="Complete")
    except Exception as e:
        return f"{logs}\nUpload failed: {str(e)}", "Error"

    return logs, "Success"


# Build the Gradio 6 Interface
with gr.Blocks(title=" ") as demo:

    with gr.Row():
        with gr.Column():
            repo_input = gr.Textbox(label="Hugging Face Repo ID", placeholder="lainlives/qwen2-test1")
            config_input = gr.Textbox(label="Config File Path", value="./config.yml")
            output_input = gr.Textbox(label="Output Directory", value="./output")
            start_btn = gr.Button("Start Process", variant="primary")

        with gr.Column():
            status_out = gr.Label(label="Status")
            log_out = gr.Textbox(label="Execution Logs", lines=10, interactive=False)

    start_btn.click(
        fn=run_merge_and_upload,
        inputs=[repo_input, config_input, output_input],
        outputs=[log_out, status_out]
    )

if __name__ == "__main__":
    install_git_repo_editable("https://github.com/lainlives/mergekit-qwen2", "mk-qwen2")
    demo.launch()