import urllib.parse

import pandas as pd


def get_image_name(url: str) -> str:
    # example url: http://upload.wikimedia.org/wikipedia/commons/6/60/MiG-9_jno_rvb.jpg
    # should return "6/60/MiG-9_jno_rvb.jpg"
    parsed = urllib.parse.urlparse(url)
    image_path = parsed.path.lstrip("/")
    if "wikipedia/commons/" in image_path:
        image_path = image_path.split("wikipedia/commons/")[-1]
    return image_path


def get_subcategories(category_tree_df: pd.DataFrame, title: str):
    category_map = {
    row['title']: row['subcategories']
    for _, row in category_tree_df.iterrows()
}
    pages_map = {
        row['title']: row['subpages']
        for _, row in category_tree_df.iterrows()
    }


    # perform graph search to get all subcategories
    subcategories = set()
    subpages = set()

    queue = [title]
    seen = set()
    while queue:
        current = queue.pop(0)
        if current in seen:
            continue
        seen.add(current)

        for subpage in pages_map.get(current, []):
            subpages.add(subpage)

        for subcategory in category_map.get(current, []):
            if subcategory not in seen:
                queue.append(subcategory)
                subcategories.add(subcategory)

    return subcategories, subpages