import os
from io import BytesIO
from pathlib import Path
import html  # <-- added for HTML escaping

import gradio as gr
import pandas as pd
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import EntryNotFoundError

# ----- CONFIG -----

# Display names <-> internal IDs
DATASETS = {
    "DRIVE": "DRIVE",
    "CHASE_DB1": "CHASE_DB1",
    "STARE": "STARE",
    "HRF": "HRF",
}

DEFAULT_DATASET_NAME = "DRIVE"

DATASET_FILES = {
    "DRIVE": "leaderboard_DRIVE.csv",
    "CHASE_DB1": "leaderboard_CHASE_DB1.csv",
    "STARE": "leaderboard_STARE.csv",
    "HRF": "leaderboard_HRF.csv",
}

COLUMNS = [
    "model",
    "mIoU",
    "F1 score",
    "Accuracy",
    "AUC",
    "Sensitivity",
    "MCC (Matthews Correlation Coefficient)",
    "code_repo",
]

SPACE_ID = os.environ.get("SPACE_ID")
if SPACE_ID is None:
    SPACE_ID = "Zomba/Retinal-Vessel-Segmentation-Leaderboard"

API = HfApi()


def dataset_name_to_id(name: str) -> str:
    return DATASETS.get(name, list(DATASETS.values())[0])


def get_filename(dataset_name: str) -> str:
    dataset_id = dataset_name_to_id(dataset_name)
    return DATASET_FILES[dataset_id]


# ----- DATA IO: load from / save to HUB (Git) -----

def load_raw_leaderboard(dataset_name: str) -> pd.DataFrame:
    """
    Load CSV from the Space repo on the Hub.
    If it doesn't exist, return an empty DataFrame with the right columns.
    """
    filename = get_filename(dataset_name)

    try:
        # Download latest file from the Hub to a temp location
        local_path = hf_hub_download(
            repo_id=SPACE_ID,
            repo_type="space",
            filename=filename,
            local_dir="hf_cache",
            local_dir_use_symlinks=False,
        )
        df = pd.read_csv(local_path)
    except EntryNotFoundError:
        # First time: no file yet
        df = pd.DataFrame(columns=COLUMNS)

    # Ensure all columns exist in order
    for col in COLUMNS:
        if col not in df.columns:
            df[col] = None

    return df[COLUMNS]


def save_raw_leaderboard(dataset_name: str, df: pd.DataFrame) -> None:
    """
    Upload the updated CSV back to this Space repo using HfApi.upload_file.
    This overwrites the file in the repo and makes the data persistent.
    """
    filename = get_filename(dataset_name)

    csv_bytes = df.to_csv(index=False).encode("utf-8")
    fileobj = BytesIO(csv_bytes)

    API.upload_file(
        path_or_fileobj=fileobj,
        path_in_repo=filename,
        repo_id=SPACE_ID,
        repo_type="space",
    )


# ----- HTML RENDERING -----

def prepare_for_display(df: pd.DataFrame) -> str:
    """
    Sort by mIoU and render an HTML table.
    The 'model' cell becomes a clickable link using 'code_repo'.
    The 'code_repo' column itself is hidden from the table.
    """
    df = df.copy()

    # Sort by primary metric
    if "mIoU" in df.columns:
        df = df.sort_values(by="mIoU", ascending=False, na_position="last")

    # Columns to display (omit 'code_repo')
    display_cols = [
        "model",
        "mIoU",
        "F1 score",
        "Accuracy",
        "AUC",
        "Sensitivity",
        "MCC (Matthews Correlation Coefficient)",
    ]

    # Build HTML table
    html_parts = []
    html_parts.append('<div style="overflow-x:auto;">')
    html_parts.append('<table style="border-collapse:collapse; width:100%;">')

    # Header
    html_parts.append("<thead><tr>")
    for col in display_cols:
        html_parts.append(
            f'<th style="border:1px solid #ddd; padding:8px; text-align:center;">{html.escape(col)}</th>'
        )
    html_parts.append("</tr></thead>")

    # Body
    html_parts.append("<tbody>")
    for _, row in df.iterrows():
        html_parts.append("<tr>")
        for col in display_cols:
            val = row.get(col, "")

            # Model column -> clickable link using code_repo (if valid URL)
            if col == "model":
                model_name = "" if pd.isna(val) else str(val)
                url = str(row.get("code_repo") or "")
                if url.startswith("http"):
                    cell_html = f'<a href="{html.escape(url)}" target="_blank">{html.escape(model_name)}</a>'
                else:
                    cell_html = html.escape(model_name)
            else:
                if pd.isna(val):
                    cell_html = ""
                else:
                    cell_html = html.escape(str(val))

            html_parts.append(
                f'<td style="border:1px solid #ddd; padding:8px; text-align:center;">{cell_html}</td>'
            )
        html_parts.append("</tr>")
    html_parts.append("</tbody></table></div>")

    return "".join(html_parts)


def get_leaderboard(dataset_name: str) -> str:
    df = load_raw_leaderboard(dataset_name)
    return prepare_for_display(df)


# ----- SUBMIT LOGIC -----

def submit_result(
    dataset_name: str,
    model: str,
    metric1: float,
    metric2: float,
    metric3: float,
    metric4: float,
    metric5: float,
    metric6: float,
    code_repo: str,
):
    model = (model or "").strip()
    code_repo = (code_repo or "").strip()

    if not model:
        return gr.update(), "❌ Please enter a model name."

    df = load_raw_leaderboard(dataset_name)

    new_row = {
        "model": model,
        "mIoU": metric1,
        "F1 score": metric2,
        "Accuracy": metric3,
        "AUC": metric4,
        "Sensitivity": metric5,
        "MCC (Matthews Correlation Coefficient)": metric6,
        "code_repo": code_repo,
    }

    # If model already exists for this dataset, update it; else append
    mask = df["model"] == model
    if mask.any():
        df.loc[mask, COLUMNS] = [new_row[col] for col in COLUMNS]
        msg = f"✅ Updated existing entry for `{model}` in **{dataset_name}**."
    else:
        df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
        msg = f"✅ Added new entry for `{model}` to **{dataset_name}**."

    # Save (commit) to Hub
    try:
        save_raw_leaderboard(dataset_name, df)
    except Exception as e:
        # If commit fails, keep UI but show error
        return prepare_for_display(df), f"⚠️ Saved locally but failed to commit to Hub: `{e}`"

    return prepare_for_display(df), msg


def change_dataset(dataset_name: str):
    html_table = get_leaderboard(dataset_name)
    return html_table, ""


def refresh_leaderboard(dataset_name: str):
    return get_leaderboard(dataset_name)


# ----- UI -----

with gr.Blocks(title="Retinal-Vessel-Segmentation Leaderboard") as demo:
    gr.Markdown("# 🏆 Retinal-Vessel-Segmentation Leaderboard")
    gr.Markdown(
        "Use the toggle to switch between datasets.\n\n"
        "- Each dataset has its **own CSV file in the repo**.\n"
        "- Submissions are saved by **committing the CSV** back to this Space.\n"
        "- The **model name** is a clickable link to the code repository."
    )

    dataset_selector = gr.Radio(
        choices=list(DATASETS.keys()),
        value=DEFAULT_DATASET_NAME,
        label="Select dataset",
        interactive=True,
    )

    # ⬇️ Changed: use HTML instead of DataFrame
    leaderboard_html = gr.HTML(
        value=get_leaderboard(DEFAULT_DATASET_NAME),
        label="Current leaderboard",
    )

    refresh_btn = gr.Button("🔄 Refresh leaderboard")

    gr.Markdown("## 🚀 Submit / Update Result for Selected Dataset")

    with gr.Row():
        model_in = gr.Textbox(label="Model name", placeholder="my-awesome-model")

    with gr.Row():
        m1 = gr.Number(label="mIoU")
        m2 = gr.Number(label="F1 score")
        m3 = gr.Number(label="Accuracy")
    with gr.Row():
        m4 = gr.Number(label="AUC")
        m5 = gr.Number(label="Sensitivity")
        m6 = gr.Number(label="MCC (Matthews Correlation Coefficient)")

    code_repo_in = gr.Textbox(
        label="Code repository URL",
        placeholder="https://github.com/you/your-model",
    )

    submit_btn = gr.Button("💾 Submit / Update")
    submit_msg = gr.Markdown()

    # Toggle dataset -> reload table
    dataset_selector.change(
        fn=change_dataset,
        inputs=dataset_selector,
        outputs=[leaderboard_html, submit_msg],
    )

    # Refresh button
    refresh_btn.click(
        fn=refresh_leaderboard,
        inputs=dataset_selector,
        outputs=leaderboard_html,
    )

    # Submit button
    submit_btn.click(
        fn=submit_result,
        inputs=[
            dataset_selector,
            model_in,
            m1, m2, m3, m4, m5, m6,
            code_repo_in,
        ],
        outputs=[leaderboard_html, submit_msg],
    )

demo.launch()