import os
import requests
import zipfile
import io
import gradio as gr

def get_file_summary(file_info):
    return {
        "name": file_info.filename,
        "type": "binary" if file_info.file_size > 1024 * 1024 else "text",
        "size": file_info.file_size,
    }

def extract_repo_content(url):
    if "huggingface.co" not in url:
        return "Invalid URL. Please provide a valid Hugging Face URL."
    
    repo_name = url.split('/')[-2]
    repo_type = url.split('/')[-3]
    api_url = f"https://huggingface.co/api/{repo_type}/{repo_name}/tree/main"
    
    response = requests.get(api_url)
    if response.status_code != 200:
        return f"Failed to fetch repository content. Status code: {response.status_code}"
    
    repo_content = response.json()
    extracted_content = []
    headers = []
    
    for file_info in repo_content:
        file_summary = get_file_summary(file_info)
        headers.append(file_summary)
        
        if file_summary["type"] == "text" and file_summary["size"] <= 1024 * 1024:
            file_url = f"https://huggingface.co/{repo_type}/{repo_name}/resolve/main/{file_info['filename']}"
            file_response = requests.get(file_url)
            if file_response.status_code == 200:
                file_content = file_response.text
                extracted_content.append({"header": file_summary, "content": file_content})
            else:
                extracted_content.append({"header": file_summary, "content": "Failed to fetch file content."})
        else:
            extracted_content.append({"header": file_summary, "content": "File too large or binary, content not captured."})
    
    return extracted_content

def format_output(extracted_content):
    formatted_output = ""
    for file_data in extracted_content:
        formatted_output += f"### File: {file_data['header']['name']}\n"
        formatted_output += f"**Type:** {file_data['header']['type']}\n"
        formatted_output += f"**Size:** {file_data['header']['size']} bytes\n"
        formatted_output += "#### Content:\n"
        formatted_output += f"```\n{file_data['content']}\n```\n\n"
    return formatted_output

def extract_and_display(url):
    extracted_content = extract_repo_content(url)
    formatted_output = format_output(extracted_content)
    return formatted_output

app = gr.Blocks()

with app:
    gr.Markdown("# Gradio Space/Model Content Extractor")
    url_input = gr.Textbox(label="Hugging Face Space/Model URL")
    output_display = gr.Markdown()
    extract_button = gr.Button("Extract Content")
    
    extract_button.click(fn=extract_and_display, inputs=url_input, outputs=output_display)

app.launch()