import spaces
import gradio as gr
import os
from repo_utils import extract_repo_content
from model_utils import get_model_summary, install_flash_attn

# Install required package
install_flash_attn()

def format_output(extracted_content, repo_url):
    formatted_output = f"# Repository URL: {repo_url}\n\n"
    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 += f"**Created:** {file_data['header']['creation_date']}\n"
            formatted_output += f"**Modified:** {file_data['header']['modification_date']}\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):
    hf_token = os.getenv("HF_TOKEN")
    hf_user = os.getenv("SPACE_AUTHOR_NAME")
    if not hf_token or not hf_user:
        return "Error: HF_TOKEN or SPACE_AUTHOR_NAME environment variable is not set."
    extracted_content = extract_repo_content(url, hf_token, hf_user)
    formatted_output = format_output(extracted_content, url)
    return formatted_output

def handle_model_summary(model_name):
    model_summary, error_message = get_model_summary(model_name)
    if error_message:
        return error_message, ""
    return model_summary, ""

# Create Gradio App
app = gr.Blocks(theme="sudeepshouche/minimalist")

with app:
    with gr.Tab("Repository Extraction"):
        gr.Markdown("# Hugging Face Space / Model Repository Content Extractor")
        url_input = gr.Textbox(label="https:// URL of Repository", placeholder="Enter the repository URL here OR select an example below...")
        url_examples = gr.Examples(
            examples=[
                ["https://huggingface.co/spaces/big-vision/paligemma-hf"],
                ["https://huggingface.co/google/paligemma-3b-mix-224"],
                ["https://huggingface.co/microsoft/Phi-3-vision-128k-instruct"],
                ["https://huggingface.co/llava-hf/llava-v1.6-mistral-7b-hf"]
            ],
            inputs=url_input
        )
        extract_button = gr.Button("Extract Content")
        output_display = gr.Textbox(label="Extracted Repository Content", show_copy_button=True, lines=20, placeholder="Repository content will be extracted here...\n\nMetadata is captured for all files, but text content provided only for files less than 16 kb\n\n\n\nReview and search through the content here OR simply copy it for offline analysis!!. 🤖")
        extract_button.click(fn=extract_and_display, inputs=url_input, outputs=output_display)
    
    with gr.Tab("Model Explorer"):
        gr.Markdown("## Retrieve and Display Model Architecture")
        model_name_input = gr.Textbox(label="Model Name", placeholder="Enter the model name to retrieve its architecture...")
        model_output = gr.Textbox(label="Model Architecture", lines=20, placeholder="Model architecture will appear here...", show_copy_button=True)
        error_output = gr.Textbox(label="Error", lines=10, placeholder="Exceptions will appear here...", show_copy_button=True)
        model_submit_button = gr.Button("Submit")
        model_submit_button.click(fn=handle_model_summary, inputs=model_name_input, outputs=[model_output, error_output])

    with gr.Tab("Non-HF Repository Extraction"):
        gr.Markdown("# Non-Hugging Face Repository Content Extractor")
        non_hf_url_input = gr.Textbox(label="Repository URL", placeholder="Enter the non-Hugging Face repository URL here...")
        non_hf_output_display = gr.Textbox(label="Extracted Repository Content", show_copy_button=True, lines=20, placeholder="Repository content will be extracted here...")
        non_hf_extract_button = gr.Button("Extract Content")
        non_hf_extract_button.click(fn=extract_and_display, inputs=non_hf_url_input, outputs=non_hf_output_display)
        
        gr.Markdown("### Filter and Selector Options")
        file_type_filter = gr.CheckboxGroup(label="File Types", choices=[".txt", ".py", ".md", ".json", ".yaml", ".csv"])
        file_size_slider = gr.Slider(label="File Size Limit (KB)", minimum=0, maximum=1024, value=16, step=1)
        intelligent_retrieval_checkbox = gr.Checkbox(label="Enable Intelligent Retrieval")
        
        non_hf_extract_button.click(
            fn=extract_and_display,
            inputs=[non_hf_url_input, file_type_filter, file_size_slider, intelligent_retrieval_checkbox],
            outputs=non_hf_output_display
        )

if __name__ == "__main__":
    app.launch()