import requests
import json
import gradio as gr
import re


def fetch_tags(image_url):
    # Try to extract the image ID from the URL
    match = re.search(r'https://danbooru.donmai.us/posts/(\d+)', image_url)
    if not match:
        return '[ERROR]: Invalid image URL format. Please use a valid Danbooru URL.', '', '', ''

    image_id = match.group(1)  # Extract image ID from URL

    base_url = 'https://danbooru.donmai.us/posts'
    response = requests.get(f'{base_url}/{image_id}.json')

    if response.status_code != 200:
        return f'[ERROR]: {response.status_code} - Failed to retrieve data.', '', '', ''

    data = json.loads(response.text)

    # Extract required fields
    character = data.get('tag_string_character', 'N/A').replace('_', ' ')
    origin = data.get('tag_string_copyright', 'N/A').replace('_', ' ')
    tags = data.get('tag_string_general', '').replace('_', ' ')

    # Split tags into a list and join with commas
    tag_list = tags.split()  # Split on whitespace to preserve multi-word tags
    formatted_tags = ', '.join(tag_list)  # Join with comma and space

    # Prepare prompt with cleaned-up fields
    prompt = f'{character}, {origin}, {formatted_tags}'

    return character, origin, formatted_tags, prompt


# Create a Gradio interface
iface = gr.Interface(
    fn=fetch_tags,
    inputs=gr.Textbox(label="Danbooru Image URL"),
    outputs=[
        gr.Textbox(label="Character"),
        gr.Textbox(label="Origin"),
        gr.Textbox(label="Tags (comma-separated)"),
        gr.Textbox(label="Prompt (comma-separated)")
    ],
    title="Danbooru Tag Extractor",
    description="Enter a Danbooru image URL to fetch and format character, origin, tags, and prompt information."
)

# Launch the interface
iface.launch()