import json
import sys

import datasets
import gradio as gr
from apscheduler.schedulers.background import BackgroundScheduler
from huggingface_hub import snapshot_download
from loguru import logger

from envs import LOG_LEVEL

# Set the log level to INFO
logger.remove()
logger.add(sys.stdout, level=LOG_LEVEL, diagnose=False)

from app_configs import DEFAULT_SELECTIONS, THEME
from components.hf_pipelines import create_hf_pipeline_submission_interface
from components.quizbowl.bonus import BonusInterface
from components.quizbowl.tossup import TossupInterface
from components.typed_dicts import PipelineInterfaceDefaults, TossupInterfaceDefaults
from display.css_html_js import fonts_header, js_head, leaderboard_css
from display.custom_css import css_bonus, css_pipeline, css_tossup
from display.guide import BUILDING_MARKDOWN, QUICKSTART_MARKDOWN

# Constants
from envs import (
    API,
    COMPETITION_URL,
    CONTACT_EMAIL,
    DISCORD_URL,
    DOCS_REPO_URL,
    EVAL_REQUESTS_PATH,
    EVAL_RESULTS_PATH,
    GITHUB_ISSUES_URL,
    LEADERBOARD_REFRESH_INTERVAL,
    LEADERBOARD_URL,
    PLAYGROUND_DATASET_NAMES,
    QUEUE_REPO,
    REGISTRATION_URL,
    REPO_ID,
    RESULTS_REPO,
    SERVER_RESTART_INTERVAL,
)
from hf_datasets_utils import download_dataset_snapshot
from shared.workflows import factory
from shared.workflows.configs import AVAILABLE_MODELS
from shared.workflows.llms import llm_cache


def restart_space():
    llm_cache.sync_to_hf()
    API.restart_space(repo_id=REPO_ID)


download_dataset_snapshot(QUEUE_REPO, EVAL_REQUESTS_PATH)


def load_dataset(mode: str):
    if mode == "tossup":
        ds = datasets.load_dataset(PLAYGROUND_DATASET_NAMES["tossup"], split="eval")
        ds = ds.filter(lambda x: x["qid"].split("-")[2] == "1" and int(x["qid"].split("-")[3]) <= 10)
    elif mode == "bonus":
        ds = datasets.load_dataset(PLAYGROUND_DATASET_NAMES["bonus"], split="eval")
        ds = ds.filter(lambda x: x["qid"].split("-")[2] == "1" and int(x["qid"].split("-")[3]) <= 10)
    else:
        raise ValueError(f"Invalid mode: {mode}")

    return ds


def get_default_tab_id(request: gr.Request):
    logger.info(f"Request: {request}")
    tab_key_value = request.query_params.get("tab", "tossup")
    return gr.update(selected=tab_key_value)


def presave_pipeline_state(
    login_btn,
    browser_state: dict,
    tossup_pipeline_state: dict,
    tossup_output_state: dict,
    bonus_pipeline_state: dict,
    bonus_output_state: dict,
):
    browser_state.setdefault("tossup", {})
    browser_state["tossup"]["pipeline_state"] = tossup_pipeline_state
    browser_state["tossup"]["output_state"] = tossup_output_state
    browser_state.setdefault("bonus", {})
    browser_state["bonus"]["pipeline_state"] = bonus_pipeline_state
    browser_state["bonus"]["output_state"] = bonus_output_state
    logger.debug(
        f"Pipeline state before login. Login button: {login_btn}, browser state: {json.dumps(browser_state, indent=4)}"
    )
    return login_btn, browser_state


if __name__ == "__main__":
    scheduler = BackgroundScheduler()
    scheduler.add_job(restart_space, "interval", seconds=SERVER_RESTART_INTERVAL)
    scheduler.start()

    css = css_pipeline + css_tossup + css_bonus + leaderboard_css
    head = fonts_header + js_head
    tossup_ds = load_dataset("tossup")
    bonus_ds = load_dataset("bonus")
    with gr.Blocks(
        css=css,
        head=head,
        theme=THEME,
        title="QANTA 2025 Quizbowl Competition",
    ) as demo:
        browser_state = gr.BrowserState(
            {
                "tossup": {"pipeline_state": None, "output_state": None},
                "bonus": {"pipeline_state": None, "output_state": None},
            }
        )
        with gr.Row():
            with gr.Column(scale=5):
                gr.Markdown(
                    f"## 🤖 Welcome to QANTA 2025 Quizbowl Arena! &emsp;&emsp;&emsp; 👉 🏆 [Leaderboard]({LEADERBOARD_URL}) 👈"
                    "\n### 🎲 Create, play around, and submit your quizbowl agents."
                    f"<br>📋 [Register]({REGISTRATION_URL}) to participate in our [QANTA 2025 Human-AI Quizbowl Competition]({COMPETITION_URL}).",
                    elem_classes="welcome-text",
                )
            # leaderboard_btn = gr.Button("🏆 Leaderboard", elem_id="leaderboard-btn", link=LEADERBOARD_URL)
            login_btn = gr.LoginButton(scale=1)

        accent_color = "#B00000"
        gr.Markdown(
            f"👀 <span style='color:{accent_color};font-weight:bold;'>First time here?</span> Check out the [❓ Help](#help) tab for a quick introduction and "
            f"[QANTA25 Documentation]({DOCS_REPO_URL}) "
            "for detailed examples and tutorials on how to create and compete with your own QuizBowl agents. "
            f"<br>🤨 <span style='color:{accent_color};font-weight:bold;'>Facing any issues?</span> Please contact us at [{CONTACT_EMAIL}](mailto:{CONTACT_EMAIL}), or raise an issue on our [GitHub repository]({GITHUB_ISSUES_URL})."
            f" &nbsp; 🧑🏻‍💻 <span style='color:{accent_color};font-weight:bold;'>Engagements?</span> Join our [Discord server]({DISCORD_URL}) for discussions and getting help.",
            elem_classes="help-text",
        )
        with gr.Tabs() as gtab:
            with gr.Tab("🛎️ Tossup Agents", id="tossup"):
                defaults = TossupInterfaceDefaults(
                    **DEFAULT_SELECTIONS["tossup"], init_workflow=factory.create_simple_qb_tossup_workflow()
                )
                tossup_interface = TossupInterface(demo, browser_state, tossup_ds, AVAILABLE_MODELS, defaults)
            with gr.Tab("🙋🏻‍♂️ Bonus Round Agents", id="bonus"):
                defaults = PipelineInterfaceDefaults(
                    **DEFAULT_SELECTIONS["bonus"], init_workflow=factory.create_simple_qb_bonus_workflow()
                )
                bonus_interface = BonusInterface(demo, browser_state, bonus_ds, AVAILABLE_MODELS, defaults)
            # with gr.Tab("🤗 HuggingFace Pipelines", elem_id="hf-pipeline-tab", id="hf-pipeline-tab"):
            #     hf_pipeline_interface = create_hf_pipeline_submission_interface(demo)
            with gr.Tab("❓ Help", id="help"):
                with gr.Row():
                    with gr.Column():
                        gr.Markdown(QUICKSTART_MARKDOWN)
                    with gr.Column():
                        gr.Markdown(BUILDING_MARKDOWN)

        # Event Listeners

        # This is used to retrieve the pipeline state user was working on before login.
        # makes things less annoying when progress is lost due to login.
        login_btn.click(
            fn=presave_pipeline_state,
            inputs=[
                login_btn,
                browser_state,
                tossup_interface.pipeline_state,
                tossup_interface.output_state,
                bonus_interface.pipeline_state,
                bonus_interface.output_state,
            ],
            outputs=[login_btn, browser_state],
        )

        demo.queue(default_concurrency_limit=40).launch()