src/leaderboard_manager.py

"""
리더보드 관리 모듈
리더보드 데이터의 로드, 저장, 표시 준비를 담당합니다.

- 로컬 CSV: 프로젝트 루트의 data/leaderboard_results.csv
- 선택적 HF 연동:
    - repo_id: Config.FRESHQA_DATA_REPO_ID
    - token  : Config.HF_TOKEN
    - 파일명 : leaderboard_results.csv (repo 루트)
- Config.UPLOAD_LEADERBOARD_TO_HF == True 일 때만 HF를 읽고/쓴다.
"""

import os
import time
import tempfile
from typing import Optional

import pandas as pd
from huggingface_hub import HfApi, hf_hub_download

from config import Config
from src.utils import file_lock


# -------------------------
# 상수 및 설정
# -------------------------

HF_LEADERBOARD_FILENAME = "leaderboard_results.csv"          # HF dataset 내 파일명 (루트)
LOCAL_LEADERBOARD_FILENAME = "leaderboard_results.csv"      # 로컬 data 폴더 내 파일명 (기존 유지)

HF_REPO_ID = Config.FRESHQA_DATA_REPO_ID
HF_ADMIN_TOKEN = Config.HF_TOKEN
UPLOAD_LEADERBOARD_TO_HF = Config.UPLOAD_LEADERBOARD_TO_HF

hf_api = HfApi()


# -------------------------
# 경로/초기 스키마/정규화 헬퍼
# -------------------------

def _get_local_leaderboard_path() -> str:
    """프로젝트 루트 기준 로컬 리더보드 CSV 경로 반환."""
    current_dir = os.path.dirname(os.path.abspath(__file__))  # src/ 폴더
    project_root = os.path.dirname(current_dir)               # 프로젝트 루트
    return os.path.join(project_root, "data", LOCAL_LEADERBOARD_FILENAME)


def _init_empty_leaderboard_df() -> pd.DataFrame:
    """초기 빈 리더보드 스키마 DataFrame."""
    return pd.DataFrame({
        "id": [],
        "model": [],
        "description": [],
        "accuracy": [],
        "fast_changing_accuracy": [],
        "slow_changing_accuracy": [],
        "never_changing_accuracy": [],
        "acc_vp": [],
        "acc_fp": [],
        "acc_vp_one_hop": [],
        "acc_vp_two_hop": [],
        "acc_fp_one_hop": [],
        "acc_fp_two_hop": [],
        "acc_vp_old": [],
        "acc_vp_new": [],
        "acc_fp_old": [],
        "acc_fp_new": [],
        "acc_politics": [],
        "acc_sports": [],
        "acc_entertainment": [],
        "acc_weather": [],
        "acc_world": [],
        "acc_economy": [],
        "acc_society": [],
        "acc_it_science": [],
        "acc_life_culture": [],
        "acc_unknown": [],
        "total_questions": [],
        "evaluation_date": [],
        "evaluation_mode": [],
    })


def _normalize_leaderboard_df(df: pd.DataFrame) -> pd.DataFrame:
    """
    리더보드 DF를 스키마/정렬/컬럼 순서 기준에 맞춰 정규화한다.
    (기존 load_leaderboard_data의 로직을 함수로 분리)
    """
    if df is None or df.empty:
        return _init_empty_leaderboard_df()

    df = df.copy()

    # evaluation_mode가 없으면 추가
    if "evaluation_mode" not in df.columns:
        df["evaluation_mode"] = "Unknown"

    # 텍스트 컬럼 보정
    text_columns = ["model", "description"]
    for col in text_columns:
        if col not in df.columns:
            df[col] = pd.Series(dtype="object")

    # 상세 분석 컬럼 없으면 추가
    detailed_columns = [
        "acc_test", "acc_dev", "acc_vp", "acc_fp", "acc_vp_one_hop", "acc_vp_two_hop",
        "acc_fp_one_hop", "acc_fp_two_hop", "acc_vp_old", "acc_vp_new", "acc_fp_old", "acc_fp_new",
    ]
    for col in detailed_columns:
        if col not in df.columns:
            df[col] = 0.0

    # 도메인별 정확도 컬럼 없으면 추가
    domain_columns = [
        "acc_politics", "acc_sports", "acc_entertainment",
        "acc_weather", "acc_world", "acc_economy",
        "acc_society", "acc_it_science", "acc_life_culture", "acc_unknown",
    ]
    for col in domain_columns:
        if col not in df.columns:
            df[col] = 0.0

    # accuracy 기준 정렬
    if "accuracy" in df.columns and not df.empty:
        df = df.sort_values("accuracy", ascending=False).reset_index(drop=True)

    # 컬럼 순서 정렬 (rank 제외)
    column_order = [
        "id", "model", "description", "accuracy", "fast_changing_accuracy",
        "slow_changing_accuracy", "never_changing_accuracy", "acc_vp", "acc_fp",
        "acc_vp_one_hop", "acc_vp_two_hop", "acc_fp_one_hop", "acc_fp_two_hop",
        "acc_vp_old", "acc_vp_new", "acc_fp_old", "acc_fp_new",
        "acc_politics", "acc_sports", "acc_entertainment", "acc_weather",
        "acc_world", "acc_economy", "acc_society", "acc_it_science",
        "acc_life_culture", "acc_unknown", "total_questions",
        "evaluation_date", "evaluation_mode",
    ]
    available_columns = [col for col in column_order if col in df.columns]
    df = df[available_columns]

    return df


def _load_local_leaderboard_df() -> pd.DataFrame:
    """로컬 CSV에서 리더보드 로드 (없으면 빈 스키마)."""
    data_path = _get_local_leaderboard_path()
    try:
        df = pd.read_csv(data_path)
        return _normalize_leaderboard_df(df)
    except FileNotFoundError:
        return _init_empty_leaderboard_df()
    except Exception as e:
        print(f"⚠️ 로컬 리더보드 로드 실패: {e}")
        return _init_empty_leaderboard_df()


# -------------------------
# HF 연동 헬퍼
# -------------------------

def _can_use_hf() -> bool:
    """HF 연동이 가능한 상태인지 여부 (Config 기반)."""
    if not UPLOAD_LEADERBOARD_TO_HF:
        return False
    if not HF_REPO_ID or not HF_ADMIN_TOKEN:
        # 설정이 없으면 HF는 건너뜀
        return False
    return True


def _load_leaderboard_from_hf(retries: int = 3, delay: float = 1.0) -> Optional[pd.DataFrame]:
    """
    HF dataset에서 리더보드 CSV를 다운로드하여 DataFrame으로 반환.
    실패 시 None 반환. 재시도 로직 포함.
    """
    if not _can_use_hf():
        return None

    last_err: Optional[Exception] = None
    for attempt in range(1, retries + 1):
        try:
            with tempfile.TemporaryDirectory() as tmpdir:
                file_path = hf_hub_download(
                    repo_id=HF_REPO_ID,
                    filename=HF_LEADERBOARD_FILENAME,
                    repo_type="dataset",
                    local_dir=tmpdir,
                    token=HF_ADMIN_TOKEN,
                )
                df = pd.read_csv(file_path)
                return _normalize_leaderboard_df(df)
        except Exception as e:
            last_err = e
            print(f"⚠️ HF 리더보드 로드 실패 (시도 {attempt}/{retries}): {e}")
            if attempt < retries:
                time.sleep(delay)
                delay *= 2
    print("❌ HF 리더보드 로드 재시도 모두 실패")
    return None


def _save_leaderboard_to_hf(df: pd.DataFrame, retries: int = 3, delay: float = 1.0) -> bool:
    """
    HF dataset에 리더보드 CSV 업로드.
    실패 시 False 반환. 재시도 로직 포함.
    """
    if not _can_use_hf():
        return False

    df = _normalize_leaderboard_df(df)

    last_err: Optional[Exception] = None
    for attempt in range(1, retries + 1):
        try:
            with tempfile.NamedTemporaryFile(
                mode="w",
                encoding="utf-8",
                suffix=".csv",
                delete=False,
            ) as tmpfile:
                df.to_csv(tmpfile.name, index=False)
                tmp_path = tmpfile.name

            hf_api.upload_file(
                path_or_fileobj=tmp_path,
                path_in_repo=HF_LEADERBOARD_FILENAME,
                repo_id=HF_REPO_ID,
                repo_type="dataset",
                token=HF_ADMIN_TOKEN,
                commit_message="Update leaderboard results",
            )

            os.unlink(tmp_path)
            return True

        except Exception as e:
            last_err = e
            print(f"⚠️ HF 리더보드 업로드 실패 (시도 {attempt}/{retries}): {e}")
            if attempt < retries:
                time.sleep(delay)
                delay *= 2

    print(f"❌ HF 리더보드 업로드 재시도 모두 실패: {last_err}")
    return False


# -------------------------
# 공개 API: 로드 / 추가
# -------------------------

def load_leaderboard_data() -> pd.DataFrame:
    """
    리더보드 데이터 로드.

    동작 우선순위:
    1) Config.UPLOAD_LEADERBOARD_TO_HF == True && HF 설정 OK:
        - HF에서 최신 리더보드 로드 시도
        - 성공 시: 그 내용을 로컬 CSV에 덮어쓴 뒤 반환
        - 실패 시: 로컬 CSV를 사용 (없으면 빈 스키마)
    2) 그 외:
        - 로컬 CSV만 사용 (없으면 빈 스키마)
    """
    data_path = _get_local_leaderboard_path()
    lock_path = data_path + ".lock"

    # HF를 사용할 수 있는 경우에만 HF 우선 시도
    if _can_use_hf():
        with file_lock(lock_path):
            hf_df = _load_leaderboard_from_hf()
            if hf_df is not None:
                # HF가 소스 오브 트루스: 로컬 CSV도 HF 기준으로 동기화
                try:
                    os.makedirs(os.path.dirname(data_path), exist_ok=True)
                    hf_df.to_csv(data_path, index=False)
                except Exception as e:
                    print(f"⚠️ 로컬 리더보드 동기화 실패: {e}")
                return hf_df

            # HF에서 못 가져오면 로컬로 폴백
            local_df = _load_local_leaderboard_df()
            return local_df

    # HF를 사용하지 않는 경우: 로컬만
    return _load_local_leaderboard_df()


def append_to_leaderboard_data(new_data_list):
    """
    리더보드 데이터에 새로운 결과 추가 (파일 잠금 사용).

    - 항상 로컬 CSV를 업데이트
    - Config.UPLOAD_LEADERBOARD_TO_HF == True 이고 HF 설정이 유효하면,
      업데이트된 전체 DF를 HF에도 업로드 (재시도 포함).
    """
    data_path = _get_local_leaderboard_path()
    lock_path = data_path + ".lock"

    with file_lock(lock_path):
        # 1) 로컬 기존 데이터 로드
        if os.path.exists(data_path):
            try:
                existing_df = pd.read_csv(data_path)
            except Exception as e:
                print(f"⚠️ 로컬 리더보드 읽기 실패, 빈 스키마로 진행: {e}")
                existing_df = _init_empty_leaderboard_df()
        else:
            existing_df = _init_empty_leaderboard_df()

        existing_df = _normalize_leaderboard_df(existing_df)

        # 2) 새로운 데이터 추가
        new_df = pd.DataFrame(new_data_list)
        if not new_df.empty:
            new_df = _normalize_leaderboard_df(new_df)

        frames_to_concat = []
        if not existing_df.empty:
            frames_to_concat.append(existing_df)
        if not new_df.empty:
            frames_to_concat.append(new_df)

        if len(frames_to_concat) == 0:
            combined_df = existing_df.copy()
        elif len(frames_to_concat) == 1:
            combined_df = frames_to_concat[0].copy()
        else:
            combined_df = pd.concat(frames_to_concat, ignore_index=True)

        combined_df = _normalize_leaderboard_df(combined_df)

        # 3) 로컬 저장
        try:
            os.makedirs(os.path.dirname(data_path), exist_ok=True)
            combined_df.to_csv(data_path, index=False)
        except Exception as e:
            print(f"❌ 로컬 리더보드 저장 실패: {e}")

        # 4) HF에도 업로드 (옵션)
        if _can_use_hf():
            ok = _save_leaderboard_to_hf(combined_df)
            if not ok:
                print("⚠️ 리더보드 HF 업로드 실패 (로컬에는 저장됨)")

        return combined_df