import requests
import gradio as gr
import json

def get_file_summary(file_info):
    return {
        "name": file_info['path'],
        "type": "binary" if file_info['size'] > 1024 * 1024 else "text",
        "size": file_info['size'],
    }

def extract_repo_content(url):
    if "huggingface.co" not in url:
        return [{"header": {"name": "Error", "type": "error", "size": 0}, "content": "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 [{"header": {"name": "Error", "type": "error", "size": 0}, "content": f"Failed to fetch repository content. Status code: {response.status_code}"}]
    
    repo_content = response.json()
    extracted_content = []
    
    for file_info in repo_content:
        file_summary = get_file_summary(file_info)
        content = {"header": 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['path']}"
            file_response = requests.get(file_url)
            if file_response.status_code == 200:
                content["content"] = file_response.text
            else:
                content["content"] = "Failed to fetch file content."
        else:
            content["content"] = "File too large or binary, content not captured."
        
        extracted_content.append(content)
    
    return extracted_content

def format_output(extracted_content):
    formatted_output = ""
    for file_data in extracted_content:
        if isinstance(file_data, dict) and 'header' in file_data:
            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"
        else:
            formatted_output += "Error in file data format.\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()