import json
import os
import time
import shutil
import tempfile

import gradio as gr
import pandas as pd
from huggingface_hub import snapshot_download

from src.display.css_html_js import custom_css
from src.leaderboard.student_results import get_failed_submissions_by_group, get_group_submission_history, get_student_leaderboard_df
from src.submission.student_queue import queue_student_submission
from src.envs import EVAL_REQUESTS_PATH, EVAL_RESULTS_PATH, PROJ_DIR, QUEUE_REPO, RESULTS_REPO, TOKEN


demo = gr.Blocks(css=custom_css)
with demo:
    gr.HTML("<h1 align='center' id='space-title'>IMG2GPS Student Leaderboard</h1>")
    gr.Markdown(
        "Upload your PyTorch checkpoint (.pt), `model.py`, and `preprocess.py` to evaluate on the shared dataset. "
        "A separate Leaderboard tab shows best-per-alias results and stats.",
        elem_classes="markdown-text",
    )

    with gr.Tabs(elem_classes="tab-buttons") as tabs:
        with gr.TabItem("🎓 Student Submissions", elem_id="student-submissions-tab"):
            with gr.Column():
                with gr.Row():
                    group_id_tb = gr.Textbox(label="Group ID", placeholder="Your pure digit on the sheet", scale=1)
                with gr.Row():
                    alias_tb = gr.Textbox(label="Alias", placeholder="Nickname / Team name (shown on leaderboard)", scale=1)
                with gr.Row():
                    state_dict_uploader = gr.File(label="Model state dict (.pt/.pth)", file_count="single", file_types=[".pt", ".pth"])
                with gr.Row():
                    model_py_uploader = gr.File(label="Model class (model.py)", file_count="single", file_types=[".py"])
                    preproc_py_uploader = gr.File(label="Preprocessing (preprocess.py)", file_count="single", file_types=[".py"])
                with gr.Row():
                    eval_btn = gr.Button(
                        "Upload and Queue Evaluation",
                        variant="primary",
                        elem_id="submit-eval-btn",
                    )
                with gr.Row():
                    student_submit_result = gr.Markdown(elem_id="submit-result")

            def _file_path(obj):
                if obj is None:
                    return None
                if isinstance(obj, dict) and "name" in obj:
                    return obj["name"]
                if hasattr(obj, "name"):
                    return obj.name
                return str(obj)


            def _stage_uploaded_file(obj, expected_name: str, staging_dir: str, required: bool = True):
                """
                Copy a Gradio temp upload into our own temp directory before queueing.
                This prevents /tmp/gradio/... paths from disappearing during upload/commit.
                """
                src = _file_path(obj)

                if src is None:
                    if required:
                        raise FileNotFoundError(f"Missing required file: {expected_name}")
                    return None

                if not os.path.isfile(src):
                    raise FileNotFoundError(
                        f"{expected_name} was not available on the local filesystem. "
                        "Please re-upload the file and submit again."
                    )

                dst = os.path.join(staging_dir, expected_name)
                shutil.copy2(src, dst)

                if not os.path.isfile(dst):
                    raise FileNotFoundError(f"Failed to stage {expected_name}")

                if os.path.getsize(dst) == 0:
                    raise ValueError(f"{expected_name} is empty")

                return dst

            def handle_student_submit(group_id, alias, state_dict_file, model_py_file, preproc_py_file, progress=gr.Progress(track_tqdm=False)):
                group_id_norm = str(group_id).strip()
                alias_raw = alias or ""
                alias_clean = alias_raw.strip() if alias_raw.strip() else "NA"

                if not group_id_norm:
                    yield "❌ **Error:** Group ID is required."
                    return

                # Immediate visual feedback
                yield "⏳ **Processing...** Uploading files and queueing evaluation. Please wait."

                progress(0.05, desc="Validating submission...")

                try:
                    with tempfile.TemporaryDirectory(prefix="frontend_stage_submission_") as staging_dir:
                        staged_state_dict = _stage_uploaded_file(
                            state_dict_file,
                            "model.pt",
                            staging_dir,
                            required=False,
                        )
                        staged_model_py = _stage_uploaded_file(
                            model_py_file,
                            "model.py",
                            staging_dir,
                            required=True,
                        )
                        staged_preproc_py = _stage_uploaded_file(
                            preproc_py_file,
                            "preprocess.py",
                            staging_dir,
                            required=True,
                        )

                        queue_msg, ts = queue_student_submission(
                            group_id=group_id_norm,
                            alias=alias_raw,
                            state_dict_file=staged_state_dict,
                            model_py_file=staged_model_py,
                            preproc_py_file=staged_preproc_py,
                        )
                except Exception as e:
                    yield f"❌ **Error while queueing submission:** {e}"
                    return

                yield f"✅ **Submission queued!** Waiting for evaluation worker...\n\n{queue_msg}"
                progress(0.2, desc="Queued. Waiting for evaluation worker...")

                max_wait_s = 600  # 10 minutes
                poll_interval_s = 5
                start_time = time.time()
                status = "PENDING"
                last_error = ""

                folder_name = f"{group_id_norm} {alias_clean}"

                poll_count = 0
                while time.time() - start_time < max_wait_s and status.upper() in ("PENDING", "PENDING_NEW_EVAL", "RUNNING"):
                    elapsed = time.time() - start_time
                    frac = min(1.0, elapsed / max_wait_s)
                    progress(0.2 + 0.6 * frac, desc=f"Evaluating... (status: {status})")

                    # Show live status updates every few polls
                    poll_count += 1
                    elapsed_min = int(elapsed // 60)
                    elapsed_sec = int(elapsed % 60)
                    status_emoji = "🔄" if status == "RUNNING" else "⏳"
                    yield f"{status_emoji} **Evaluating...** Status: `{status}` | Elapsed: {elapsed_min}m {elapsed_sec}s\n\n{queue_msg}"

                    # Refresh queue mirror, but tolerate failures
                    try:
                        snapshot_download(
                            repo_id=QUEUE_REPO,
                            local_dir=EVAL_REQUESTS_PATH,
                            repo_type="dataset",
                            tqdm_class=None,
                            etag_timeout=30,
                            token=TOKEN,
                            allow_patterns=[f"{PROJ_DIR}/**/request.json"],
                        )
                    except Exception:
                        pass

                    request_path = os.path.join(EVAL_REQUESTS_PATH, PROJ_DIR, folder_name, ts, "request.json")
                    if os.path.isfile(request_path):
                        try:
                            with open(request_path, "r") as f:
                                data = json.load(f)
                            status = str(data.get("status", "PENDING")).upper()
                            last_error = data.get("error", "") or ""
                        except Exception:
                            pass

                    if status in ("SUCCESS", "FINISHED", "FAILED"):
                        break

                    time.sleep(poll_interval_s)

                if status in ("SUCCESS", "FINISHED"):
                    progress(1.0, desc="Evaluation complete.")
                    yield f"✅ **Success!** {queue_msg}\n\n🎉 Evaluation successful! Check the **Leaderboard** and **Submission Status** tabs."
                elif status == "FAILED":
                    progress(1.0, desc="Evaluation failed.")
                    err_text = last_error or "No error message provided by evaluator."
                    yield f"❌ **Evaluation Failed**\n\n{queue_msg}\n\n**Error:**\n```\n{err_text}\n```"
                else:
                    progress(0.95, desc="Still pending in evaluation queue.")
                    yield f"⏳ **Pending**\n\n{queue_msg}\n\nYour submission is still in the evaluation queue. Please check **Submission Status** or **Leaderboard** later."

            eval_btn.click(
                handle_student_submit,
                [group_id_tb, alias_tb, state_dict_uploader, model_py_uploader, preproc_py_uploader],
                [student_submit_result],
            )


        with gr.TabItem("📊 Leaderboard", elem_id="student-leaderboard-tab"):
            # Ensure latest results are cached locally at startup
            try:
                snapshot_download(
                    repo_id=RESULTS_REPO,
                    local_dir=EVAL_RESULTS_PATH,
                    repo_type="dataset",
                    tqdm_class=None,
                    etag_timeout=30,
                    token=TOKEN,
                    allow_patterns=[f"{PROJ_DIR}/student_attempts/*.json"],
                )
            except Exception:
                pass

            # img_val
            df_val, stats_val = get_student_leaderboard_df(dataset_name="img_val")
            stats_val_md_text = (
                f"[img_val] Avg distance (m) mean ± std: {stats_val['avg_distance_m_mean']:.2f} ± {stats_val['avg_distance_m_std']:.2f} | "
                f"Avg inference (ms) mean ± std: {stats_val['avg_ms_mean']:.2f} ± {stats_val['avg_ms_std']:.2f}"
                if df_val is not None and not df_val.empty
                else "[img_val] No submissions yet."
            )
            with gr.Accordion("img_val", open=True):
                stats_md_val = gr.Markdown(value=stats_val_md_text)
                leaderboard_df_val = gr.components.Dataframe(
                    value=df_val if df_val is not None else pd.DataFrame(),
                    row_count=5,
                )
            refresh_btn = gr.Button("Refresh Leaderboards")

            def refresh_leaderboards():
                try:
                    snapshot_download(
                        repo_id=RESULTS_REPO,
                        local_dir=EVAL_RESULTS_PATH,
                        repo_type="dataset",
                        tqdm_class=None,
                        etag_timeout=30,
                        token=TOKEN,
                        allow_patterns=[f"{PROJ_DIR}/student_attempts/*.json"],
                    )
                except Exception:
                    pass
                df1, s1 = get_student_leaderboard_df(dataset_name="img_val")
                if df1 is not None and not df1.empty:
                    s1_text = (
                        f"[img_val] Avg distance (m) mean ± std: {s1['avg_distance_m_mean']:.2f} ± {s1['avg_distance_m_std']:.2f} | "
                        f"Avg inference (ms) mean ± std: {s1['avg_ms_mean']:.2f} ± {s1['avg_ms_std']:.2f}"
                    )
                else:
                    s1_text = "[img_val] No submissions yet."
                return df1, s1_text

            refresh_btn.click(
                refresh_leaderboards,
                inputs=[],
                outputs=[leaderboard_df_val, stats_md_val],
            )

        with gr.TabItem("📥 Submission Status", elem_id="student-status-tab"):
            # Guidelines warning banner
            gr.Markdown(
                """
## ⚠️ Important: Read Before Submitting

**Please follow the submission guidelines carefully before submitting your code!**

Common reasons for failed submissions include:
- Missing or incorrectly named files (`model.py`, `preprocess.py`)
- Incompatible model architecture or missing dependencies
- Runtime errors in your code (check your imports and function signatures)
- Incorrect output format from your model

**Review the submission requirements** to avoid wasting your submission attempts.

---
                """,
                elem_classes="markdown-text",
            )

            # Shared Group ID input at the top
            gr.Markdown("### 🔍 Check Your Submission Status\nEnter your Group ID to view both successful and failed submissions:")
            with gr.Row():
                status_group_id_tb = gr.Textbox(
                    label="Group ID",
                    placeholder="Your pure digit on the sheet",
                    scale=1,
                )
            with gr.Row():
                check_status_btn = gr.Button("Check Submission Status")

            # Section 1: Successful Submissions
            with gr.Accordion("✅ Successful Submissions", open=True):
                status_result_md = gr.Markdown()
                status_perf_df = gr.components.Dataframe(row_count=1)

            # Section 2: Failed Submissions
            with gr.Accordion("❌ Failed Submissions", open=True):
                gr.Markdown(
                    "*If you have failed submissions, review the error messages below and fix the issues before resubmitting.*"
                )
                failed_result_md = gr.Markdown()
                failed_submissions_df = gr.components.Dataframe(row_count=1)

            def check_submission_status(group_id: str):
                group_id = (group_id or "").strip()
                if not group_id:
                    return (
                        "Please enter a valid Group ID.",
                        pd.DataFrame(),
                        "",
                        pd.DataFrame(),
                    )

                # Ensure we have the latest data locally
                try:
                    snapshot_download(
                        repo_id=RESULTS_REPO,
                        local_dir=EVAL_RESULTS_PATH,
                        repo_type="dataset",
                        tqdm_class=None,
                        etag_timeout=30,
                        token=TOKEN,
                        allow_patterns=[f"{PROJ_DIR}/student_attempts/*.json"],
                    )
                    snapshot_download(
                        repo_id=QUEUE_REPO,
                        local_dir=EVAL_REQUESTS_PATH,
                        repo_type="dataset",
                        tqdm_class=None,
                        etag_timeout=30,
                        token=TOKEN,
                        allow_patterns=[f"{PROJ_DIR}/**/request.json"],
                    )
                except Exception:
                    pass

                # Get successful submissions
                submissions_df = get_group_submission_history(group_id)
                if submissions_df is None or submissions_df.empty:
                    success_md = f"No successful submissions found for Group ID `{group_id}`."
                    success_df = pd.DataFrame()
                else:
                    success_md = f"**Total successful submissions**: {len(submissions_df)}\n\nMost recent submission appears first."
                    history_rows = []
                    for _, sub in submissions_df.iterrows():
                        history_rows.append({
                            "Alias": sub.get("alias", ""),
                            "Timestamp": sub.get("timestamp", ""),
                            "Dataset": sub.get("dataset", ""),
                            "Avg distance (m)": sub.get("avg_distance_m"),
                            "Avg infer (ms)": sub.get("avg_infer_ms"),
                            "Total infer (s)": sub.get("total_infer_s"),
                        })
                    success_df = pd.DataFrame(history_rows)

                # Get failed submissions
                failed_df = get_failed_submissions_by_group(group_id)
                if failed_df is None or failed_df.empty:
                    failed_md = f"No failed submissions found for Group ID `{group_id}`. 🎉"
                    failed_df = pd.DataFrame()
                else:
                    failed_md = f"**Total failed submissions**: {len(failed_df)}\n\n⚠️ Please review the error messages and fix the issues before resubmitting."

                return success_md, success_df, failed_md, failed_df

            check_status_btn.click(
                check_submission_status,
                inputs=[status_group_id_tb],
                outputs=[status_result_md, status_perf_df, failed_result_md, failed_submissions_df],
            )

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