import requests
import gradio as gr

HF_API_BASE = "https://huggingface.co/api/models"
SEARCH_TERMS = [
    "mj", "michaeljackson", "michael_jackson", "michael-jackson", "mjackson",
    "m_jackson", "m-jackson", "kingofpop", "king_of_pop", "king-of-pop"
]

def get_models(user: str):
    url = f"{HF_API_BASE}?search={user}"
    headers = {"User-Agent": "Mozilla/5.0"}
    try:
        response = requests.get(url, headers=headers)
        response.raise_for_status()
        return response.json() if response.text else []
    except requests.exceptions.RequestException as e:
        print(f"Error fetching models: {e}")
        return []

def check_zip_files(model_id: str):
    url = f"https://huggingface.co/api/models/{model_id}/tree/main"
    headers = {"User-Agent": "Mozilla/5.0"}
    zip_links = []
    try:
        response = requests.get(url, headers=headers)
        response.raise_for_status()
        if response.text:
            files = response.json()
            for file in files:
                file_path = file.get('path', '').lower()
                if file.get('type') == 'file' and file_path.endswith('.zip'):
                    if any(term in file_path for term in SEARCH_TERMS):
                        zip_links.append(f"https://huggingface.co/{model_id}/resolve/main/{file['path']}?download=true")
    except requests.exceptions.RequestException as e:
        print(f"Error accessing model {model_id}: {e}")
    return zip_links

def search_huggingface_zips(user: str):
    models = get_models(user)
    all_links = []
    for model in models:
        model_id = model.get('modelId', '')
        if model_id:
            all_links.extend(check_zip_files(model_id))
    return "\n".join(all_links) if all_links else "No ZIP files found."

demo = gr.Interface(
    fn=search_huggingface_zips,
    inputs=gr.Textbox(label="Hugging Face Username"),
    outputs=gr.Textbox(label="Download Links"),
    title="Michael Jackson Models Finder",
    description="Enter a Hugging Face username to search for MJ-related ZIP files in their models.",
)

demo.launch()