import gradio as gr
from PIL import Image
from PIL.PngImagePlugin import PngInfo
import requests
from bs4 import BeautifulSoup
import os
import subprocess
from huggingface_hub import ModelCard, SpaceCard

# PNG Info Tool
def png_info(file):
    img = Image.open(file)
    info = img.info
    return str(info)

# Web Info Tool (with HTML Info)
def web_info(url):
    # Fetch the webpage
    response = requests.get(url)
    soup = BeautifulSoup(response.text, 'html.parser')
    
    # Extract basic info
    title = soup.title.string if soup.title else "No Title"
    meta_description = soup.find("meta", attrs={"name": "description"})
    description = meta_description["content"] if meta_description else "No Description"
    
    # Extract additional HTML info
    meta_tags = soup.find_all("meta")
    headers = {f"h{i}": len(soup.find_all(f"h{i}")) for i in range(1, 7)}
    links = [a["href"] for a in soup.find_all("a", href=True)]
    
    # Format the output
    basic_info = f"Title: {title}\nDescription: {description}"
    html_info = f"Meta Tags: {len(meta_tags)}\nHeaders: {headers}\nLinks: {len(links)}"
    return f"{basic_info}\n\nHTML Info:\n{html_info}"

# View Source Info Tool
def view_source_info(url):
    if url.startswith("view-source:"):
        url = url.replace("view-source:", "").strip()
    if not url.startswith(("http://", "https://")):
        url = "https://" + url
    response = requests.get(url)
    return response.text

# Document Info Tool
def document_info(file):
    file_size = os.path.getsize(file)
    file_type = os.path.splitext(file)[1]
    return f"File Type: {file_type}\nFile Size: {file_size} bytes"

# Video Info Tool
def video_info(file):
    # Display the video
    video_display = file
    
    # Extract video metadata using ffprobe
    result = subprocess.run(
        ["ffprobe", "-v", "error", "-show_format", "-show_streams", file],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
    )
    
    # Return the video display and metadata
    return video_display, result.stdout

# Model Info Tool
def model_info(model_id):
    # Fetch model card from Hugging Face Hub
    card = ModelCard.load(model_id)
    return card.content  # Return the model card content (markdown)

# Space Info Tool
def space_info(space_id):
    # Fetch space card from Hugging Face Hub
    card = SpaceCard.load(space_id)
    return card.content  # Return the space card content (markdown)

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## Information Extraction Tools")
    
    with gr.Tab("PNG Info"):
        png_input = gr.File(label="Upload PNG File")
        png_output = gr.Textbox(label="PNG Info")
        png_button = gr.Button("Extract PNG Info")
        png_button.click(png_info, inputs=png_input, outputs=png_output)
    
    with gr.Tab("Web Info"):
        web_input = gr.Textbox(label="Enter URL")
        web_output = gr.Textbox(label="Web Info")
        web_button = gr.Button("Extract Web Info")
        web_button.click(web_info, inputs=web_input, outputs=web_output)
    
    with gr.Tab("View Source Info"):
        source_input = gr.Textbox(label="Enter URL (with or without 'view-source:')", placeholder="e.g., https://example.com or view-source:example.com")
        source_output = gr.Textbox(label="HTML Source Code", lines=20)
        source_button = gr.Button("View Source")
        source_button.click(view_source_info, inputs=source_input, outputs=source_output)
    
    with gr.Tab("Document Info"):
        doc_input = gr.File(label="Upload Document")
        doc_output = gr.Textbox(label="Document Info")
        doc_button = gr.Button("Extract Document Info")
        doc_button.click(document_info, inputs=doc_input, outputs=doc_output)
    
    with gr.Tab("Video Info"):
        video_input = gr.Video(label="Upload Video")
        video_output = gr.Video(label="Video Preview")
        metadata_output = gr.Textbox(label="Video Metadata", lines=10)
        video_button = gr.Button("Extract Video Info")
        video_button.click(video_info, inputs=video_input, outputs=[video_output, metadata_output])
    
    with gr.Tab("Model Info"):
        model_input = gr.Textbox(label="Enter Model ID", placeholder="e.g., bert-base-uncased")
        model_output = gr.Markdown(label="Model Info")
        model_button = gr.Button("Fetch Model Info")
        model_button.click(model_info, inputs=model_input, outputs=model_output)
    
    with gr.Tab("Space Info"):
        space_input = gr.Textbox(label="Enter Space ID", placeholder="e.g., gradio/hello-world")
        space_output = gr.Markdown(label="Space Info")
        space_button = gr.Button("Fetch Space Info")
        space_button.click(space_info, inputs=space_input, outputs=space_output)

# Launch the Gradio app
demo.launch()