# import streamlit as st
# st.write("hello")


import os
import json
import uuid
import time
import tempfile
from typing import Any, Dict, List, Tuple

import streamlit as st
import pandas as pd
from PIL import Image
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import HfHubHTTPError
import logging
import sys

# Configure logging to write to stdout (so HF captures it)
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO,
    handlers=[logging.StreamHandler(sys.stdout)]
)

logger = logging.getLogger(__name__)

# =========================
# CONFIG
# =========================

logger.info(">> Application Started !")
                           
st.set_page_config(
    page_title="AI Benchmark Arena",
    page_icon="🏆",
    layout="wide",
    initial_sidebar_state="expanded",
)

logger.info(">> Accessing tokens")
#Set this to the private dataset repo that acts as the "database"
DB_REPO_ID = os.getenv("DB_REPO_ID", "VizWiz-Challenges/submissions-db")
DB_REPO_TYPE = "dataset"
# print(DB_REPO_ID)


logger.info(">> Loading submission tokens")
# This must exist as a Space Secret in the PUBLIC UI Space
SUBMISSIONS_TOKEN = os.getenv("SUBMISSIONS_TOKEN")
logger.info(">> All tokens loaded")

# Phase config (copied from EvalAI config intent)
PHASES = [
    {"label": "Dev (qeury-dev2024)", "codename": "test-dev2024"},
    {"label": "Standard (query-standard2024)", "codename": "test-standard2024"},
    {"label": "Challenge (query-challenge2024)", "codename": "test-challenge2024"},
]

CHALLENGE_TYPES = ["Object Detection", "Instance Segmentation"]

# Leaderboard columns from your EvalAI yaml
LEADERBOARD_METRICS = ["bbox_mAP", "bbox_AP50", "segm_mAP", "segm_AP50"]
DEFAULT_SORT_METRIC = "segm_AP50"

# =========================
# HELPERS
# =========================

def _require_token() -> None:
    if not SUBMISSIONS_TOKEN:
        st.error(
            "Missing SUBMISSIONS_TOKEN. Add it in Space Settings → Secrets "
            "(token must have read/write access ONLY to the private DB dataset repo)."
        )
        st.stop()


def _validate_submission_json(obj: Any) -> Tuple[bool, str]:
    """
    Validates the submission format from your instructions:
    - Top-level must be a list
    - Each item must be a dict containing:
      image_id (int), score (number), category_id (int), area (number),
      bbox ([x,y,w,h]), segmentation (list)
    """
    if not isinstance(obj, list):
        return False, "Submission must be a JSON list of annotations."

    required_keys = {"image_id", "score", "category_id", "area", "bbox", "segmentation"}

    for i, ann in enumerate(obj):
        if not isinstance(ann, dict):
            return False, f"Annotation at index {i} must be an object/dict."

        missing = required_keys - set(ann.keys())
        if missing:
            return False, f"Annotation at index {i} missing keys: {sorted(list(missing))}"

        # Basic type checks
        if not isinstance(ann["image_id"], int):
            return False, f"image_id at index {i} must be an integer."
        if not isinstance(ann["category_id"], int):
            return False, f"category_id at index {i} must be an integer."

        if not isinstance(ann["score"], (int, float)):
            return False, f"score at index {i} must be a number."

        if not isinstance(ann["area"], (int, float)):
            return False, f"area at index {i} must be a number."

        bbox = ann["bbox"]
        if not (isinstance(bbox, list) and len(bbox) == 4 and all(isinstance(x, (int, float)) for x in bbox)):
            return False, f"bbox at index {i} must be a list of 4 numbers: [x, y, w, h]."

        segm = ann["segmentation"]
        if not isinstance(segm, list):
            return False, f"segmentation at index {i} must be a list."

    return True, "OK"


def _upload_json(api: HfApi, data: Dict[str, Any] | List[Any], path_in_repo: str) -> None:
    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tmp:
        json.dump(data, tmp, ensure_ascii=False)
        tmp_path = tmp.name

    try:
        api.upload_file(
            path_or_fileobj=tmp_path,
            path_in_repo=path_in_repo,
            repo_id=DB_REPO_ID,
            repo_type=DB_REPO_TYPE,
            token=SUBMISSIONS_TOKEN,
            commit_message=f"Add {path_in_repo}",
        )
    finally:
        try:
            os.remove(tmp_path)
        except OSError:
            pass


def _create_submission_record(
    *,
    pred: List[Dict[str, Any]],
    team: str,
    model_name: str,
    phase_codename: str,
    challenge_type: str,
    original_filename: str,
) -> str:
    """
    Writes pred/meta/status to the private DB dataset repo.
    Returns submission_id.
    """
    _require_token()
    api = HfApi()

    submission_id = str(uuid.uuid4())
    ts = int(time.time())

    meta = {
        "submission_id": submission_id,
        "team": team.strip(),
        "model": model_name.strip(),
        "phase_codename": phase_codename,
        "challenge_type": challenge_type,
        "timestamp": ts,
        "original_filename": original_filename,
    }

    status = {"state": "queued", "timestamp": ts}

    base = f"submissions/{submission_id}"
    _upload_json(api, pred, f"{base}/pred.json")
    _upload_json(api, meta, f"{base}/meta.json")
    _upload_json(api, status, f"{base}/status.json")

    return submission_id


def _download_leaderboard_jsonl() -> str | None:
    """
    Downloads leaderboard.jsonl from the DB repo.
    Returns local path or None if missing.
    """
    _require_token()
    try:
        return hf_hub_download(
            repo_id=DB_REPO_ID,
            repo_type=DB_REPO_TYPE,
            filename="leaderboard.jsonl",
            token=SUBMISSIONS_TOKEN,
        )
    except HfHubHTTPError as e:
        # Most common: 404 when file doesn't exist yet
        if "404" in str(e):
            return None
        raise


def _load_leaderboard_df() -> pd.DataFrame:
    path = _download_leaderboard_jsonl()
    if path is None:
        return pd.DataFrame(columns=["team", "model", "phase_codename", *LEADERBOARD_METRICS, "timestamp"])

    rows = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError:
                # Skip malformed lines rather than crashing the UI
                continue

    if not rows:
        return pd.DataFrame(columns=["team", "model", "phase_codename", *LEADERBOARD_METRICS, "timestamp"])

    df = pd.DataFrame(rows)

    # Ensure columns exist
    for col in ["team", "model", "phase_codename", "timestamp", *LEADERBOARD_METRICS]:
        if col not in df.columns:
            df[col] = None

    # Sort descending by default metric
    if DEFAULT_SORT_METRIC in df.columns:
        df = df.sort_values(by=DEFAULT_SORT_METRIC, ascending=False, kind="mergesort")

    return df


# =========================
# UI
# =========================
#2 FUNCTIONS WERE HERE, RENDER_OVERVIEW AND RENDER EVAL

def page_submit():
    st.header("🚀 Submit your Predictions")

    col1, col2 = st.columns([2, 1])

    with col2:
        st.subheader("Submission Info")
        team = st.text_input("Team / Display Name", value=st.session_state.get("team", ""))
        model_name = st.text_input("Model Name", value=st.session_state.get("model_name", ""))

        phase_label = st.selectbox("Phase", [p["label"] for p in PHASES])
        phase_codename = next(p["codename"] for p in PHASES if p["label"] == phase_label)

        challenge_type = st.radio("Challenge type", CHALLENGE_TYPES, horizontal=False)

        st.session_state["team"] = team
        st.session_state["model_name"] = model_name

        st.caption("Your submission will be queued for evaluation. Scores appear on the leaderboard after processing.")

    with col1:
        st.subheader("Upload Submission File")
        uploaded_file = st.file_uploader("Choose a JSON file", type=["json"])

        if uploaded_file is None:
            st.info("Upload a JSON file that contains a list of annotations.")
            return

        # Parse JSON
        try:
            raw = uploaded_file.getvalue().decode("utf-8")
            pred_obj = json.loads(raw)
        except Exception:
            st.error("Could not parse JSON. Please upload a valid JSON file.")
            return

        ok, msg = _validate_submission_json(pred_obj)
        if not ok:
            st.error(f"Invalid submission format: {msg}")
            return

        st.success("Submission file looks valid ✅")

        submit_clicked = st.button("Submit (Queue for Evaluation)", type="primary")

        if submit_clicked:
            if not team.strip():
                st.error("Please enter Team / Display Name.")
                return
            if not model_name.strip():
                st.error("Please enter Model Name.")
                return

            with st.spinner("Uploading submission to the private database repo..."):
                try:
                    submission_id = _create_submission_record(
                        pred=pred_obj,
                        team=team,
                        model_name=model_name,
                        phase_codename=phase_codename,
                        challenge_type=challenge_type,
                        original_filename=uploaded_file.name,
                    )
                except Exception as e:
                    st.error(f"Upload failed: {e}")
                    return

            st.balloons()
            st.success("Submission queued successfully!")
            st.code(f"Submission ID: {submission_id}")

            st.info("Next: the private evaluator will score your submission and update the leaderboard.")


def page_leaderboard():
    st.header("🏆 Leaderboard")
    st.write(f"Ranked by **{DEFAULT_SORT_METRIC}** (descending).")

    with st.spinner("Loading leaderboard from private database repo..."):
        try:
            df = _load_leaderboard_df()
        except Exception as e:
            st.error(f"Could not load leaderboard: {e}")
            return

    if df.empty:
        st.info("No scored submissions yet. Submit a model to get started!")
        return

    # Add Rank column
    df_display = df.copy()
    df_display.insert(0, "Rank", range(1, len(df_display) + 1))

    # Optional: pretty timestamp
    if "timestamp" in df_display.columns:
        df_display["timestamp"] = pd.to_datetime(df_display["timestamp"], unit="s", errors="coerce")

    st.dataframe(
        df_display,
        column_config={
            "Rank": st.column_config.Column("Rank", width="small"),
            "team": "Team",
            "model": "Model",
            "phase_codename": "Phase",
            "bbox_mAP": st.column_config.NumberColumn("bbox_mAP", format="%.4f"),
            "bbox_AP50": st.column_config.NumberColumn("bbox_AP50", format="%.4f"),
            "segm_mAP": st.column_config.NumberColumn("segm_mAP", format="%.4f"),
            "segm_AP50": st.column_config.NumberColumn("segm_AP50", format="%.4f"),
            "timestamp": st.column_config.DatetimeColumn("Scored at", format="D MMM YYYY, h:mm a"),
        },
        use_container_width=True,
        hide_index=True,
    )


def main():

    # st.write("✅ App booted and rendering UI")
    
    st.sidebar.title("AI Benchmark Arena 🏆")

    # Warn early if DB repo isn't configured
    if DB_REPO_ID.startswith("NidhiS09/"):
        st.sidebar.warning("Set DB_REPO_ID env var or hardcode your private DB dataset repo id in main.py.")

    menu = ["Submit Model", "Leaderboard"]
    choice = st.sidebar.radio("Navigation", menu)

    st.sidebar.markdown("---")
    st.sidebar.caption("This Space queues submissions to a private DB repo and reads leaderboard results from it.")

    render_overview()
    render_eval_details()
    st.markdown("---")

    if choice == "Submit Model":
        page_submit()
    else:
        page_leaderboard()

def render_overview():
    with st.expander("ℹ️ Overview of the AI Benchmark Arena"):
        st.markdown(
            """
        **Note:** This Hugging Face Space queues submissions for evaluation and persists results in a private database repo.
            """
        )
        # Keep this optional so missing image doesn't crash the Space
        try:
            overview_image = Image.open("src/overview_image.png").resize((600, 600))
            st.image(overview_image, caption="Example of an object localization task")
        except Exception:
            st.info("Overview image not found at src/overview_image.png (optional).")


def render_eval_details():
    with st.expander("📐 How is the Score Calculated?"):
        st.markdown(
            """
        Your submission is evaluated offline by a private evaluator against hidden ground-truth annotations.
        The leaderboard reports:

        - bbox_mAP
        - bbox_AP50
        - segm_mAP
        - segm_AP50 (default ranking)

        Raw submissions are kept private; only scores and metadata are shown.
            """
        )

main()
# if __name__ == "__main__":
#     main()