# app.py

import gradio as gr
import json
import os
import shutil
import subprocess
import tempfile
from datetime import datetime
from typing import Dict, Any
import uuid

from evaluate import evaluate_submission


# =========================================================
# Config
# =========================================================

GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
GITHUB_USERNAME = os.getenv("GITHUB_USERNAME", "YOUR_GITHUB_USERNAME")
GITHUB_REPO = os.getenv("GITHUB_REPO", "YOUR_LEADERBOARD_REPO")
GITHUB_BRANCH = os.getenv("GITHUB_BRANCH", "main")

ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD", "admin123")

# GitHub 仓库 URL
GITHUB_REPO_URL = f"https://{GITHUB_TOKEN}@github.com/{GITHUB_USERNAME}/{GITHUB_REPO}.git" if GITHUB_TOKEN else ""


# =========================================================
# GitHub 操作函数
# =========================================================

def clone_repo():
    """克隆 GitHub 仓库到临时目录"""
    if not GITHUB_TOKEN:
        raise ValueError("GitHub token not configured")
    
    temp_dir = tempfile.mkdtemp()
    result = subprocess.run(
        ["git", "clone", "--depth", "1", GITHUB_REPO_URL, temp_dir],
        capture_output=True, text=True
    )
    if result.returncode != 0:
        shutil.rmtree(temp_dir, ignore_errors=True)
        raise Exception(f"Clone failed: {result.stderr}")
    
    return temp_dir


def push_repo(repo_path, commit_message):
    """推送更改到 GitHub"""
    try:
        # 配置 Git
        subprocess.run(["git", "config", "user.email", "leaderboard@bot.com"], cwd=repo_path, capture_output=True)
        subprocess.run(["git", "config", "user.name", "Leaderboard Bot"], cwd=repo_path, capture_output=True)
        
        # 添加所有更改
        subprocess.run(["git", "add", "."], cwd=repo_path, capture_output=True)
        
        # 检查是否有更改
        status = subprocess.run(["git", "status", "--porcelain"], cwd=repo_path, capture_output=True, text=True)
        if not status.stdout.strip():
            return "ℹ️ No changes to commit"
        
        # 提交
        subprocess.run(["git", "commit", "-m", commit_message], cwd=repo_path, capture_output=True)
        
        # 推送
        push_result = subprocess.run(
            ["git", "push", "origin", GITHUB_BRANCH],
            cwd=repo_path,
            capture_output=True,
            text=True
        )
        
        if push_result.returncode != 0:
            return f"❌ Push failed: {push_result.stderr}"
        
        return "✅ Successfully pushed to GitHub"
    except Exception as e:
        return f"❌ Git operation failed: {str(e)}"


def load_json_from_repo(repo_path, file_path):
    """从克隆的仓库中加载 JSON 文件"""
    full_path = os.path.join(repo_path, file_path)
    if not os.path.exists(full_path):
        return [] if file_path.endswith('.json') else {}
    with open(full_path, "r", encoding="utf-8") as f:
        return json.load(f)


def save_json_to_repo(repo_path, file_path, data):
    """保存 JSON 文件到克隆的仓库"""
    full_path = os.path.join(repo_path, file_path)
    os.makedirs(os.path.dirname(full_path), exist_ok=True)
    with open(full_path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)


# =========================================================
# Submit
# =========================================================

def submit_result(
    prediction_dyadic,
    prediction_multiparty,
    method_name,
    backbone,
    category,
    organization,
    paper_url,
    code_url,
    contact_email
):
    print("=" * 50)
    print("开始处理提交...")
    
    # =====================================================
    # Required Checks
    # =====================================================
    
    if prediction_dyadic is None:
        return "❌ prediction_dyadic.json is required."
    
    if prediction_multiparty is None:
        return "❌ prediction_multiparty.json is required."
    
    if not method_name.strip():
        return "❌ Method Name is required."
    
    if not backbone.strip():
        return "❌ Backbone Model is required."
    
    if not category:
        return "❌ Category is required."
    
    # =====================================================
    # Filename Checks
    # =====================================================
    
    if os.path.basename(prediction_dyadic.name) != "prediction_dyadic.json":
        return "❌ File must be named prediction_dyadic.json"
    
    if os.path.basename(prediction_multiparty.name) != "prediction_multiparty.json":
        return "❌ File must be named prediction_multiparty.json"
    
    repo_path = None
    try:
        # 1. 克隆 GitHub 仓库
        print("克隆 GitHub 仓库...")
        repo_path = clone_repo()
        
        # 2. 加载现有数据
        pending_submissions = load_json_from_repo(repo_path, "pending_submissions.json")
        pending_leaderboard = load_json_from_repo(repo_path, "pending_leaderboard.json")
        
        # 3. 生成 submission_id
        existing_ids = [item.get("submission_id", 0) for item in pending_submissions]
        submission_id = max(existing_ids) + 1 if existing_ids else 1
        print(f"Submission ID: {submission_id}")
        
        # 4. 保存上传的预测文件到 submissions 文件夹
        submission_folder = os.path.join(repo_path, "submissions", str(submission_id))
        os.makedirs(submission_folder, exist_ok=True)
        
        dyadic_path = os.path.join(submission_folder, "prediction_dyadic.json")
        multiparty_path = os.path.join(submission_folder, "prediction_multiparty.json")
        
        shutil.copy(prediction_dyadic.name, dyadic_path)
        shutil.copy(prediction_multiparty.name, multiparty_path)
        print("预测文件已保存到仓库")
        
        # 5. 评估
        print("开始评估...")
        evaluation_result = evaluate_submission(
            dyadic_path, multiparty_path, method_name, backbone,
            category, organization, paper_url, code_url
        )
        print("评估完成")
        
        pending_leaderboard_entry = evaluation_result.get("pending_leaderboard_entry", {})
        dataset_scores = evaluation_result.get("dataset_scores", {})
        
        if not pending_leaderboard_entry:
            return "❌ 评估结果为空，请检查 prediction 文件格式"
        
        # 6. 添加元数据
        pending_leaderboard_entry["submission_id"] = submission_id
        pending_leaderboard_entry["status"] = "pending"
        pending_leaderboard_entry["submission_time"] = datetime.now().isoformat()
        
        # 7. 更新 pending_leaderboard
        pending_leaderboard.append(pending_leaderboard_entry)
        save_json_to_repo(repo_path, "pending_leaderboard.json", pending_leaderboard)
        
        # 8. 更新 pending_submissions
        pending_submission_entry = {
            "submission_id": submission_id,
            "method_name": method_name,
            "category": category,
            "backbone": backbone,
            "dataset_scores": dataset_scores,
            "organization": organization,
            "paper_url": paper_url,
            "code_url": code_url,
            "contact_email": contact_email,
            "status": "pending",
            "submission_time": datetime.now().isoformat()
        }
        pending_submissions.append(pending_submission_entry)
        save_json_to_repo(repo_path, "pending_submissions.json", pending_submissions)
        
        # 9. 推送到 GitHub
        print("推送到 GitHub...")
        push_result = push_repo(repo_path, f"Add submission {submission_id}: {method_name}")
        print(push_result)
        
        # 10. 生成返回信息
        result_msg = f"""
✅ Submission Successful

Submission ID: {submission_id}
Method: {method_name}
Backbone: {backbone}
Category: {category}

Scores (D&M Dataset):
  LLM: {dataset_scores.get('D&M', {}).get('LLM', 0):.2f}
  P: {dataset_scores.get('D&M', {}).get('P', 0):.2f}
  R: {dataset_scores.get('D&M', {}).get('R', 0):.2f}
  F1: {dataset_scores.get('D&M', {}).get('F1', 0):.2f}
  BLEU-1: {dataset_scores.get('D&M', {}).get('BLEU-1', 0):.2f}

Status: Pending Admin Approval
{push_result}
"""
        return result_msg
        
    except Exception as e:
        return f"❌ Submission failed: {str(e)}"
    
    finally:
        # 清理临时目录
        if repo_path and os.path.exists(repo_path):
            shutil.rmtree(repo_path, ignore_errors=True)
            print("清理临时目录完成")


# =========================================================
# Admin Functions
# =========================================================

def admin_login(password):
    if password == ADMIN_PASSWORD:
        return (gr.update(visible=True), "✅ Login Successful")
    return (gr.update(visible=False), "❌ Wrong Password")


def get_pending_submissions():
    """从 GitHub 获取待审批的提交列表"""
    repo_path = None
    try:
        repo_path = clone_repo()
        pending_submissions = load_json_from_repo(repo_path, "pending_submissions.json")
        
        rows = []
        for item in pending_submissions:
            if item.get("status") == "pending":
                dataset_scores = item.get("dataset_scores", {})
                dm_scores = dataset_scores.get("D&M", {})
                
                rows.append([
                    item.get("submission_id", ""),
                    item.get("method_name", ""),
                    item.get("category", ""),
                    item.get("backbone", ""),
                    f"{dm_scores.get('LLM', 0):.2f}",
                    f"{dm_scores.get('P', 0):.2f}",
                    f"{dm_scores.get('R', 0):.2f}",
                    f"{dm_scores.get('F1', 0):.2f}",
                    f"{dm_scores.get('BLEU-1', 0):.2f}"
                ])
        return rows
    except Exception as e:
        print(f"Error: {e}")
        return []
    finally:
        if repo_path and os.path.exists(repo_path):
            shutil.rmtree(repo_path, ignore_errors=True)


def approve_submission(submission_id):
    """批准提交，将结果移动到 leaderboard.json"""
    try:
        submission_id = int(submission_id)
    except ValueError:
        return "❌ Invalid Submission ID"
    
    repo_path = None
    try:
        # 1. 克隆仓库
        repo_path = clone_repo()
        
        # 2. 加载数据
        pending_submissions = load_json_from_repo(repo_path, "pending_submissions.json")
        pending_leaderboard = load_json_from_repo(repo_path, "pending_leaderboard.json")
        leaderboard = load_json_from_repo(repo_path, "components/data/leaderboard.json")
        
        # 3. 查找并验证提交
        approved_entry = None
        pending_sub_entry = None
        
        for item in pending_leaderboard:
            if item.get("submission_id") == submission_id and item.get("status") == "pending":
                approved_entry = item
                break
        
        for item in pending_submissions:
            if item.get("submission_id") == submission_id and item.get("status") == "pending":
                pending_sub_entry = item
                break
        
        if approved_entry is None:
            return "❌ Submission Not Found"
        
        # 4. 准备 leaderboard 条目（移除 submission_id 和 status）
        leaderboard_entry = {
            "method": approved_entry.get("method", {}),
            "results": approved_entry.get("results", {})
        }
        
        # 5. 检查是否已存在相同的方法
        existing_index = None
        for i, entry in enumerate(leaderboard):
            if (entry.get("method", {}).get("name") == leaderboard_entry["method"].get("name") and
                entry.get("method", {}).get("backbone") == leaderboard_entry["method"].get("backbone")):
                existing_index = i
                break
        
        # 6. 更新或添加
        if existing_index is not None:
            leaderboard[existing_index] = leaderboard_entry
            action = "updated"
        else:
            leaderboard.append(leaderboard_entry)
            action = "added"
        
        # 7. 更新状态
        approved_entry["status"] = "approved"
        if pending_sub_entry:
            pending_sub_entry["status"] = "approved"
        
        # 8. 保存文件
        save_json_to_repo(repo_path, "components/data/leaderboard.json", leaderboard)
        save_json_to_repo(repo_path, "pending_leaderboard.json", pending_leaderboard)
        save_json_to_repo(repo_path, "pending_submissions.json", pending_submissions)
        
        # 9. 推送到 GitHub
        push_result = push_repo(repo_path, f"Approve submission {submission_id}: {approved_entry['method']['name']} ({action})")
        
        return f"""
✅ Approved: {submission_id}

Method: {approved_entry.get('method', {}).get('name')}
Backbone: {approved_entry.get('method', {}).get('backbone')}
Action: {action}

{push_result}
"""
    except Exception as e:
        return f"❌ Approval failed: {str(e)}"
    finally:
        if repo_path and os.path.exists(repo_path):
            shutil.rmtree(repo_path, ignore_errors=True)


def reject_submission(submission_id):
    """拒绝提交"""
    try:
        submission_id = int(submission_id)
    except ValueError:
        return "❌ Invalid Submission ID"
    
    repo_path = None
    try:
        # 1. 克隆仓库
        repo_path = clone_repo()
        
        # 2. 加载数据
        pending_submissions = load_json_from_repo(repo_path, "pending_submissions.json")
        pending_leaderboard = load_json_from_repo(repo_path, "pending_leaderboard.json")
        
        # 3. 更新状态
        for item in pending_leaderboard:
            if item.get("submission_id") == submission_id and item.get("status") == "pending":
                item["status"] = "rejected"
        
        for item in pending_submissions:
            if item.get("submission_id") == submission_id and item.get("status") == "pending":
                item["status"] = "rejected"
        
        # 4. 保存文件
        save_json_to_repo(repo_path, "pending_leaderboard.json", pending_leaderboard)
        save_json_to_repo(repo_path, "pending_submissions.json", pending_submissions)
        
        # 5. 推送到 GitHub
        push_result = push_repo(repo_path, f"Reject submission {submission_id}")
        
        return f"❌ Rejected: {submission_id}\n\n{push_result}"
    except Exception as e:
        return f"❌ Rejection failed: {str(e)}"
    finally:
        if repo_path and os.path.exists(repo_path):
            shutil.rmtree(repo_path, ignore_errors=True)


# =========================================================
# UI
# =========================================================

with gr.Blocks(title="H2HMem Leaderboard") as demo:
    
    gr.Markdown("# 🏆 H2HMem Leaderboard")
    gr.Markdown("""
    ### Submission Instructions
    
    Please upload:
    
    - **`prediction_dyadic.json`** - Predictions for dyadic conversations
    - **`prediction_multiparty.json`** - Predictions for multiparty conversations
    
    Both files are required for evaluation.
    """)
    
    with gr.Tabs():
        
        # =================================================
        # Submit Tab
        # =================================================
        
        with gr.Tab("📤 Submit"):
            with gr.Row():
                with gr.Column():
                    prediction_dyadic = gr.File(label="prediction_dyadic.json *", file_types=[".json"])
                    prediction_multiparty = gr.File(label="prediction_multiparty.json *", file_types=[".json"])
                    method_name = gr.Textbox(label="Method Name *", placeholder="e.g., A-Mem, MuRAG, etc.")
                    backbone = gr.Textbox(label="Backbone Model *", placeholder="e.g., GPT-4, LLaMA-3, etc.")
                    category = gr.Dropdown(
                        label="Category *",
                        choices=["Text-based", "Multimodal"],
                        value="Text-based"
                    )
                    organization = gr.Textbox(label="Organization (Optional)", placeholder="Your institution/company")
                    paper_url = gr.Textbox(label="Paper URL (Optional)", placeholder="https://...")
                    code_url = gr.Textbox(label="Code URL (Optional)", placeholder="https://github.com/...")
                    contact_email = gr.Textbox(label="Contact Email (Optional)", placeholder="email@example.com")
                    submit_btn = gr.Button("🚀 Submit", variant="primary")
                    submit_output = gr.Textbox(label="Status", lines=20)
                    
                    submit_btn.click(
                        fn=submit_result,
                        inputs=[
                            prediction_dyadic,
                            prediction_multiparty,
                            method_name,
                            backbone,
                            category,
                            organization,
                            paper_url,
                            code_url,
                            contact_email
                        ],
                        outputs=submit_output
                    )
        
        # =================================================
        # Admin Tab
        # =================================================
        
        with gr.Tab("🔐 Admin"):
            admin_password = gr.Textbox(label="Admin Password", type="password")
            login_btn = gr.Button("Login")
            login_status = gr.Textbox(label="Login Status")
            admin_panel = gr.Column(visible=False)
            
            with admin_panel:
                gr.Markdown("### Pending Submissions")
                refresh_btn = gr.Button("🔄 Refresh Pending")
                pending_table = gr.Dataframe(
                    headers=[
                        "ID", "Method", "Category", "Backbone",
                        "LLM", "P", "R", "F1", "BLEU-1"
                    ],
                    interactive=False
                )
                submission_id_box = gr.Textbox(label="Submission ID", placeholder="Enter submission ID")
                
                with gr.Row():
                    approve_btn = gr.Button("✅ Approve", variant="primary")
                    reject_btn = gr.Button("❌ Reject", variant="stop")
                
                admin_output = gr.Textbox(label="Admin Status", lines=15)
                
                refresh_btn.click(fn=get_pending_submissions, outputs=pending_table)
                approve_btn.click(fn=approve_submission, inputs=submission_id_box, outputs=admin_output)
                reject_btn.click(fn=reject_submission, inputs=submission_id_box, outputs=admin_output)
                
                # 初始加载
                pending_table.value = get_pending_submissions()
            
            login_btn.click(
                fn=admin_login,
                inputs=admin_password,
                outputs=[admin_panel, login_status]
            )


demo.launch(ssr_mode=False)