import os
import shutil
import gradio as gr
from huggingface_hub import HfApi, snapshot_download

def fetch_files(source_repo, hf_token):
    if not hf_token.strip() or not source_repo.strip():
        yield gr.update(), gr.update(), "Error: Token and Source Repo required.", []
        return
        
    api = HfApi(token=hf_token.strip())
    try:
        yield gr.update(), gr.update(), f"Fetching files from '{source_repo}'...", []
        files = api.list_repo_files(repo_id=source_repo.strip(), repo_type="model")
        files.sort()
        
        yield (
            gr.update(visible=True), 
            gr.update(choices=files, value=files), 
            f"Success: Fetched {len(files)} files. Uncheck the ones to exclude.",
            files
        )
    except Exception as e:
        yield gr.update(visible=False), gr.update(), f"Error: {str(e)}", []


def clone_gated_model(source_repo, target_repo, hf_token, selected_files, all_files):
    if not target_repo.strip():
        yield "Error: Target repo required."
        return
        
    api = HfApi(token=hf_token.strip())
    local_dir = "./tmp_clone_dir"
    
    # Files the user unchecked will be passed to ignore_patterns
    files_to_ignore = list(set(all_files) - set(selected_files))
    
    try:
        yield f"1. Creating private target repo '{target_repo}'..."
        api.create_repo(repo_id=target_repo.strip(), repo_type="model", private=True, exist_ok=True)
        
        yield f"2. Downloading selected files from '{source_repo}' to Space disk..."
        snapshot_download(
            repo_id=source_repo.strip(),
            repo_type="model",
            local_dir=local_dir,
            token=hf_token.strip(),
            ignore_patterns=files_to_ignore if files_to_ignore else None
        )
        
        yield f"3. Uploading clean files to '{target_repo}'..."
        api.upload_folder(
            folder_path=local_dir,
            repo_id=target_repo.strip(),
            repo_type="model",
            commit_message="Initial commit: Cloned gated model"
        )
        
        yield f"Success! Model cloned. View it here: https://huggingface.co/{target_repo.strip()}"
        
    except Exception as e:
        yield f"An error occurred: {str(e)}"
        
    finally:
        if os.path.exists(local_dir):
            shutil.rmtree(local_dir, ignore_errors=True)


with gr.Blocks(title="Gated Model Cloner") as demo:
    gr.Markdown("# 📑 Clean Gated Model Cloner")
    gr.Markdown("Downloads a gated model to the Space's local disk, then uploads a clean copy to your account.")
    
    with gr.Row():
        hf_token = gr.Textbox(label="HF Token (Write)", type="password")
    
    with gr.Row():
        source_repo = gr.Textbox(label="Source Gated Repo", placeholder="e.g., meta-llama/Llama-3.2-1B")
        fetch_btn = gr.Button("1. Fetch Files", variant="secondary")
        
    all_files_state = gr.State([])
    
    with gr.Group(visible=False) as selection_group:
        file_checkboxes = gr.CheckboxGroup(label="Select files to keep")
        target_repo = gr.Textbox(label="Target Repo", placeholder="e.g., username/clean-llama")
        clone_btn = gr.Button("2. Clone Selected Files", variant="primary")
        
    output_logs = gr.Textbox(label="Status Logs", interactive=False)
    
    fetch_btn.click(
        fn=fetch_files,
        inputs=[source_repo, hf_token],
        outputs=[selection_group, file_checkboxes, output_logs, all_files_state]
    )
    
    clone_btn.click(
        fn=clone_gated_model,
        inputs=[source_repo, target_repo, hf_token, file_checkboxes, all_files_state],
        outputs=output_logs
    )

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Soft())