components/utils.py

import logging
from datetime import datetime, timezone
import pathlib
import tempfile
import json
from pathlib import Path
import httpx


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

from about import API, SUBMISSIONS_REPO, RESULTS_REPO, TOKEN, REGISTRATION_REPO
from components.registration.config import TEAMS_FILE_NAME, TEAM_COLUMNS
from evaluation import evaluate_problem

logger = logging.getLogger(__name__)


def make_user_clickable(name):
    link = f"https://huggingface.co/{name}"
    return f'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">{name}</a>'


def make_boundary_clickable(filename):
    link = f"https://huggingface.co/datasets/proxima-fusion/constellaration-bench-results/blob/main/{filename}"
    return f'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">link</a>'


def read_result_from_hub(filename):
    local_path = hf_hub_download(
        repo_id=RESULTS_REPO,
        repo_type="dataset",
        filename=filename,
    )
    return local_path


def read_submission_from_hub(filename):
    local_path = hf_hub_download(
        repo_id=SUBMISSIONS_REPO,
        repo_type="dataset",
        filename=filename,
    )
    return local_path


def write_results(record, result):
    record.update(result)
    record["result_filename"] = (
        record["submission_filename"].rstrip(".json") + "_results.json"
    )
    record["evaluated"] = True

    record["objectives"] = json.dumps(record.get("objectives", []))
    record["feasibilities"] = json.dumps(record.get("feasibility", []))

    if "objective" not in record.keys():
        record["objective"] = 0.0
        record["minimize_objective"] = True
        record["feasibility"] = sum(record["feasibility"]) / len(record["feasibility"])

    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tmp:
        json.dump(record, tmp)
        tmp.flush()
        tmp_name = tmp.name

    API.upload_file(
        path_or_fileobj=tmp_name,
        path_in_repo=record["result_filename"],
        repo_id=RESULTS_REPO,
        repo_type="dataset",
        commit_message=f"Add result data for {record['result_filename']}",
    )

    pathlib.Path(tmp_name).unlink()
    return


def get_user(profile: gr.OAuthProfile | None) -> str:
    if profile is None:
        return "Please login to submit a boundary for evaluation."
    return profile.username


def evaluate_boundary(filename):
    local_path = read_submission_from_hub(filename)
    with Path(local_path).open("r") as f:
        raw = f.read()
        data_dict = json.loads(raw)

    try:
        result = evaluate_problem(local_path)
    except Exception as e:
        raise gr.Error(
            f"Evaluation failed: {e}. No results written to results dataset."
        )

    write_results(data_dict, result)
    return


def show_output_box(message):
    return gr.update(value=message, visible=True)


def push_data_to_dataset(df: pd.DataFrame, repo_id: str, filename: str) -> None:
    """Upload a DataFrame as a UTF-8 CSV file to the HF dataset repository."""
    api = HfApi(token=TOKEN)
    csv_bytes = df.to_csv(index=False).encode("utf-8")
    api.upload_file(
        path_or_fileobj=csv_bytes,
        path_in_repo=filename,
        repo_id=repo_id,
        repo_type="dataset",
        commit_message=f"upsert {filename} – {datetime.now(timezone.utc).isoformat()}",
    )


def load_data_from_dataset(
    repo_id: str, filename: str, columns: list[str]
) -> pd.DataFrame:
    """
    Download a CSV file from the HF dataset and return a DataFrame.

    :param repo_id:     The HF dataset repository ID.
    :param filename:    The name of the CSV file to load.
    :return:            Returns the dataset if the file exists. Otherwise, an empty DataFrame with the correct columns.
    """
    try:
        path = hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            repo_type="dataset",
            token=TOKEN,
        )
        df = pd.read_csv(path, dtype=str).fillna("")
        for col in columns:
            if col not in df.columns:
                df[col] = ""
        return df[columns]
    except EntryNotFoundError:
        return pd.DataFrame(columns=columns)
    except Exception as exc:
        logger.error("Failed to load %s from %s: %s", filename, repo_id, exc)
        raise


def get_team(teams_df: pd.DataFrame, team_name: str) -> pd.Series | None:
    """Get the team record from teams_df corresponding to given team_name."""
    rows = teams_df[teams_df["team_name"] == team_name]
    return rows.iloc[0] if not rows.empty else None


def check_team_has_leader(team_name: str) -> bool:
    """Return True if the team exists and has a designated leader."""
    try:
        teams_df = load_data_from_dataset(
            REGISTRATION_REPO, TEAMS_FILE_NAME, TEAM_COLUMNS
        )
    except Exception:
        return False
    team = get_team(teams_df, team_name)
    if team is None:
        return False
    return bool((team.get("leader_username") or "").strip())


def prefill_user_info(
    oauth_profile: gradio.OAuthProfile | None,
    oauth_token: gradio.OAuthToken | None,
):
    empty = lambda: gr.update(value="", placeholder="Log in with HuggingFace first")
    if oauth_profile is None or oauth_token is None:
        return empty(), empty(), empty()

    try:
        resp = httpx.get(
            "https://huggingface.co/oauth/userinfo",
            headers={"Authorization": f"Bearer {oauth_token.token}"},
            timeout=5,
        )
        resp.raise_for_status()
        email = resp.json().get("email", "")
    except Exception:
        email = ''

    return (
        gr.update(value=oauth_profile.username),
        gr.update(value=oauth_profile.name),
        gr.update(value=email),
    )