import streamlit as st
import os

from app.backend.constant import (
    LEADERBOARD_ICON_MAP, get_display_name,
    DOMAIN_GROUPS, LANGUAGE_GROUPS,
)

LEADERBOARD_MAP = {}
PAGE_SECTIONS = []


def _page_toml(name, page_name):
    icon = LEADERBOARD_ICON_MAP.get(page_name, "")
    return f"""
[[pages]]
path = "app/ui/pages/{name}.py"
name = "{page_name}"
icon = "{icon}"
"""


def _section_toml(section_name):
    return f"""
[[pages]]
name = "{section_name}"
icon = ""
is_section = true
"""


def init_leaderboard():
    data_engine = st.session_state["data_engine"]
    leaderboard_map = {}
    leaderboard_change = False
    page_change = False

    text_root = []
    domain_pages = []
    language_pages = []
    multimodal_root = []
    multimodal_sub = []
    sort_id = 0

    for dataset in data_engine.datasets:
        sort_id += 1
        leaderboard = dataset["leaderboard"]
        name = dataset["name"]

        leaderboard_section = f"{leaderboard.capitalize()} Leaderboard"
        if leaderboard_section not in leaderboard_map:
            leaderboard_map[leaderboard_section] = []

        is_root = name.lower() == leaderboard.lower()
        leaderboard_map[leaderboard_section].append((name, 0 if is_root else sort_id))

        page_name = leaderboard_section if is_root else get_display_name(name)
        entry = _page_toml(name, page_name)

        if leaderboard.lower() == "text":
            if is_root:
                text_root.append(entry)
            elif name in DOMAIN_GROUPS:
                domain_pages.append(entry)
            elif name in LANGUAGE_GROUPS:
                language_pages.append(entry)
        elif leaderboard.lower() == "multimodal":
            if is_root:
                multimodal_root.append(entry)
            else:
                multimodal_sub.append(entry)

    for k, v in leaderboard_map.items():
        v.sort(key=lambda x: x[1])

    page_sections = []
    page_sections.extend(text_root)
    if domain_pages:
        page_sections.append(_section_toml("Domain"))
        page_sections.extend(domain_pages)
    if language_pages:
        page_sections.append(_section_toml("Languages"))
        page_sections.extend(language_pages)
    if multimodal_root or multimodal_sub:
        page_sections.append(_section_toml("Multimodal"))
        page_sections.extend(multimodal_root)
        page_sections.extend(multimodal_sub)

    if leaderboard_map != LEADERBOARD_MAP:
        LEADERBOARD_MAP.update(leaderboard_map)
        leaderboard_change = True
    if page_sections != PAGE_SECTIONS:
        PAGE_SECTIONS.clear()
        PAGE_SECTIONS.extend(page_sections)
        page_change = True

    return leaderboard_change, page_change


def init_pages(leaderboard_change, page_change):
    if leaderboard_change:
        with open("app/ui/pages/data_page.py", "r", encoding="utf-8") as f:
            data_page = f.read()
        for leaderboard, group_names in LEADERBOARD_MAP.items():
            for group_name in group_names:
                path = os.path.join("app/ui/pages", f"{group_name[0]}.py")
                with open(path, "w", encoding="utf-8") as f:
                    f.write(data_page.replace("$group_name$", group_name[0]))
    if page_change:
        with open("app/ui/pages_sections.toml", "w", encoding="utf-8") as f:
            f.write("\n".join(PAGE_SECTIONS))


if __name__ == '__main__':
    init_leaderboard()
    init_pages(True, True)
    print("\n".join(PAGE_SECTIONS))