import gradio
import re
import requests 

def reprocess_general_tags(general):
    tags = general.split(" ")
    tags = [tag for tag in tags if "censor" not in tag or "convenient_censoring" in tag]

    possible_person_tags = ["1boy", "2boys", "3boys", "4boys", "5boys", "6+boys", "1girl", "2girls", "3girls", "4girls", "5girls", "6+girls"]

    person_tags = [tag for tag in tags if tag in possible_person_tags]
    tags_without_person_tags = [tag for tag in tags if tag not in possible_person_tags]

    return (" ".join(person_tags), " ".join(tags_without_person_tags))

def process_tags(artist, character, copyright, general):
    txt = ""
    if artist:
        txt += artist + ","

    if character:
        txt += character + ","

    if copyright:
        txt += copyright + ","

    if general:
        person_tags, general = reprocess_general_tags(general)
        txt += general + ","
        txt = person_tags + "," + txt

    txt = txt.replace(" ", ",")
    txt = txt.replace("_", " ")

    return txt


def get_tags(url):
    reg = re.compile(r'\d{4,}')
    post_id = reg.findall(url)[0]

    if "danbooru" in url:
        url = f"https://danbooru.donmai.us/posts/{post_id}.json"

        response = requests.get(url)
        response = response.json()
        artist = response["tag_string_artist"]
        character = response["tag_string_character"]
        general = response["tag_string_general"]
        copyright = response["tag_string_copyright"]

        return process_tags(artist, character, copyright, general)

    elif "gelbooru" in url:
        url = f"https://gelbooru.com/index.php?page=dapi&q=index&json=1&s=post&id={post_id}"
        response = requests.get(url)
        response = response.json()

        general = response["post"][0]["tags"]

        return process_tags("", "", "", general)

    else:
        return "Invalid url"


iface = gradio.Interface(fn=get_tags, inputs="text", outputs="text")
iface.launch(server_port=7865)