cli.py

#!/usr/bin/env python3
import hashlib
import importlib.metadata
import json
import os
import re
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from datetime import datetime, timezone
from io import BytesIO

import click
import datasets
import httpx
from litellm import model_cost as litellm_model_cost
from litellm import register_model

from agenteval.leaderboard.schema_generator import check_submissions_against_readme

from .cli_utils import AliasedChoice, generate_choice_help
from .config import load_suite_config
from .io import atomic_write_file
from .leaderboard.models import LeaderboardSubmission
from .leaderboard.upload import (
    compress_model_usages,
    sanitize_path_component,
    upload_folder_to_hf,
)
from .models import EvalConfig, SubmissionMetadata, TaskResults
from .score import process_eval_logs
from .summary import compute_summary_statistics

HF_URL_PATTERN = r"^hf://(?:datasets/)?(?P<repo_id>[^/]+/[^/]+)/(?P<path>.*)$"
EVAL_CONFIG_FILENAME = "eval_config.json"
SCORES_FILENAME = "scores.json"
SUMMARY_FILENAME = "summary_stats.json"
SUBMISSION_METADATA_FILENAME = "submission.json"
SUMMARIES_PREFIX = "summaries"
OPENNESS_MAPPING = {
    "c": "Closed",
    "api": "API Available",
    "os": "Open Source",
    "ow": "Open Source + Open Weights",
}
TOOL_MAPPING = {
    "s": "Standard",
    "css": "Custom with Standard Search",
    "c": "Fully Custom",
}


@dataclass
class RepoPathsOfInterest:
    repo_id: str
    relative_paths: list[str]

    @staticmethod
    def from_urls(urls: list[str]):
        repo_ids = set()
        paths = []

        for url in urls:
            # validates submission_url format "hf://<repo_id>/<path>"
            repo_id, path = parse_hf_url(url)
            repo_ids.add(repo_id)
            paths.append(path)

        if len(repo_ids) > 1:
            raise ValueError("All URLs must reference the same repo")

        repo_id_to_use = repo_ids.pop()

        return RepoPathsOfInterest(
            repo_id=repo_id_to_use,
            relative_paths=paths,
        )


def parse_hf_url(url: str) -> tuple[str, str]:
    hf_url_match = re.match(HF_URL_PATTERN, url)
    if not hf_url_match:
        click.echo(
            f"Invalid URL: {url}. " "Expected format: hf://<repo_id>/<submission_path>"
        )
        sys.exit(1)

    return hf_url_match.group("repo_id"), hf_url_match.group("path")


def verify_git_reproducibility() -> None:
    try:
        # Get current commit SHA and origin
        sha_result = subprocess.run(
            ["git", "rev-parse", "--short", "HEAD"],
            capture_output=True,
            text=True,
            check=True,
        )
        origin_result = subprocess.run(
            ["git", "remote", "get-url", "origin"],
            capture_output=True,
            text=True,
            check=True,
        )
        sha = sha_result.stdout.strip() if sha_result.returncode == 0 else None
        origin = origin_result.stdout.strip() if origin_result.returncode == 0 else None

        # Check for dirty working directory
        git_dirty = (
            subprocess.run(
                ["git", "diff", "--quiet", "--exit-code"],
                capture_output=True,
                check=False,
            ).returncode
            != 0
        )

        # Warn about untracked (non-ignored) files
        untracked_result = subprocess.run(
            ["git", "ls-files", "--others", "--exclude-standard"],
            capture_output=True,
            text=True,
            check=True,
        )
        untracked_files = untracked_result.stdout.strip().splitlines()
        if untracked_files:
            click.echo(
                f"Warning: Untracked files present: {', '.join(untracked_files)}. "
                "For reproducibility, please add, ignore, or remove these files."
            )

        # Abort if worktree is dirty
        if git_dirty:
            raise click.ClickException(
                f"Git working directory contains uncommitted changes. "
                f"For reproducibility, Inspect will save: origin={origin}, sha={sha}. "
                "Please commit your changes or use --ignore-git to bypass this check (not recommended)."
            )

        # Check if commit exists on remote
        if sha:
            remote_exists = subprocess.run(
                ["git", "branch", "-r", "--contains", sha],
                capture_output=True,
                text=True,
                check=True,
            ).stdout.strip()
            if not remote_exists:
                raise click.ClickException(
                    f"Commit {sha} not found on remote '{origin}'. Others won't be able to "
                    "access this code version. Please push your changes or use --ignore-git "
                    "to bypass this check (not recommended)."
                )
    except (subprocess.SubprocessError, FileNotFoundError) as e:
        if isinstance(e, click.ClickException):
            raise
        raise click.ClickException(
            f"Unable to verify git status for reproducibility: {e}. "
            "Use --ignore-git to bypass this check if git is not available."
        )


def prep_litellm_cost_map():
    if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP") != "True":
        raise click.ClickException(
            f'Please set the LITELLM_LOCAL_MODEL_COST_MAP env variable to "True" before scoring.'
        )

    current_model_cost_keys = set(litellm_model_cost.keys())

    # Can't you just point register_model() at a URL?
    # Yes, but it won't actually use the url if LITELLM_LOCAL_MODEL_COST_MAP
    # is set, and that should be set (so we can avoid pulling costs on the fly).
    # So we'll load the cost file we want ourselves from the URL, and pass its info in as a dict.
    # This snippet is mostly lifted from
    # https://github.com/BerriAI/litellm/blob/b9621c760d3355e06dd17ec89b9eb6776755392e/litellm/litellm_core_utils/get_model_cost_map.py#L16
    # See the Development.md before changing.
    desired_model_costs_url = "https://raw.githubusercontent.com/BerriAI/litellm/eb66daeef740947c0326826817cf68fb56a8b931/litellm/model_prices_and_context_window_backup.json"
    response = httpx.get(desired_model_costs_url, timeout=5)
    response.raise_for_status()
    desired_model_costs = response.json()

    # try to check that we aren't getting info that's not also in or overridden by
    # the cost file we're pointing at
    desired_model_costs_keys = set(desired_model_costs.keys())
    in_current_not_in_desired = current_model_cost_keys - desired_model_costs_keys
    if len(in_current_not_in_desired) > 0:
        click.echo(
            f"WARNING: Info for {in_current_not_in_desired} is available but not from the specified cost map!"
        )

    register_model(model_cost=desired_model_costs)

    h = hashlib.sha256()
    h.update(json.dumps(litellm_model_cost, sort_keys=True).encode())
    model_cost_hash = h.hexdigest()
    # This is mostly informational... I think it's the case that having
    # a different hash here doesn't necessarily mean computed cost info
    # is incompatible.
    click.echo(f"Model costs hash {model_cost_hash}.")

    # Between this and the version of the file we pass to register_model()
    # I think we can reconstruct the model costs used.
    litellm_version = importlib.metadata.version("litellm")
    click.echo(f"litellm version: {litellm_version}")


@click.group()
def cli():
    pass


@click.command(
    name="score",
    help="Score a directory of evaluation logs. Can be a local directory or a HuggingFace URL.",
)
@click.argument(
    "log_dir",
    type=str,
)
def score_command(
    log_dir: str,
):
    # so that we know what model costs we're using to score
    # more details in the Development.md
    prep_litellm_cost_map()

    hf_url_match = re.match(HF_URL_PATTERN, log_dir)
    temp_dir: tempfile.TemporaryDirectory | None = None
    if hf_url_match is not None:
        # Download the logs from HF URL
        from huggingface_hub import snapshot_download

        repo_id = hf_url_match.group("repo_id")
        submission_path = hf_url_match.group("path")
        temp_dir = tempfile.TemporaryDirectory()
        download_dir = snapshot_download(
            repo_id=repo_id,
            repo_type="dataset",
            allow_patterns=f"{submission_path}/*",
            local_dir=temp_dir.__enter__(),
        )
        log_dir = os.path.join(download_dir, submission_path)

    if not os.path.exists(log_dir) or not os.path.isdir(log_dir):
        click.echo(f"No directory named {log_dir}")
        sys.exit(1)

    click.echo(f"Processing logs in {log_dir}")
    with open(os.path.join(log_dir, EVAL_CONFIG_FILENAME), "r", encoding="utf-8") as f:
        eval_config = EvalConfig.model_validate_json(f.read())

    log_processing_outcome = process_eval_logs(
        log_dir,
        reference_tasks=eval_config.suite_config.get_tasks(eval_config.split),
    )

    if log_processing_outcome.errors:
        click.echo("Errors processing logs")
        for error in log_processing_outcome.errors:
            click.echo(f"  - {error}")
        sys.exit(1)

    task_results = TaskResults(results=log_processing_outcome.results)

    # Warn if multiple evaluation specs present
    if len(task_results.agent_specs) > 1:
        click.echo(
            f"Warning: Found {len(task_results.agent_specs)} different agent configurations. "
            "Use a single solver + model config per log directory to measure a single "
            "agent's performance across tasks."
        )
    if len(task_results.code_specs) > 1:
        click.echo(
            f"Warning: Found {len(task_results.code_specs)} different code versions "
            "(revision/packages). This may indicate mixed evaluation runs from "
            "different code states."
        )

        # Warn if user-specified task arguments are present

    if task_results.tasks_with_args:
        click.echo(
            f"Warning: User-specified task arguments found for tasks: {', '.join(task_results.tasks_with_args)}. "
            "For fair comparison, do not override the task arg defaults."
        )

    # Warn about any missing tasks
    missing_tasks = eval_config.task_names - task_results.task_names
    if missing_tasks:
        click.echo(f"Warning: Missing tasks in result set: {', '.join(missing_tasks)}")

    # Persist summary
    stats = compute_summary_statistics(
        eval_config.suite_config,
        eval_config.split,
        task_results.results or [],
    )

    if hf_url_match is None:
        # Persist scores
        scores_path = os.path.join(log_dir, SCORES_FILENAME)
        atomic_write_file(scores_path, task_results.model_dump_json(indent=2))
        click.echo(f"Wrote scores to {scores_path}")

        # Persist summary
        summary_path = os.path.join(log_dir, SUMMARY_FILENAME)
        atomic_write_file(
            summary_path, json.dumps(stats.model_dump(mode="json"), indent=2)
        )
        click.echo(f"Wrote summary scores to {summary_path}")

        if temp_dir is not None:
            temp_dir.__exit__(None, None, None)
    else:
        from huggingface_hub import HfApi

        hf_api = HfApi()
        path_in_repo = f"{SUMMARIES_PREFIX}/{submission_path}/{SCORES_FILENAME}"
        hf_api.upload_file(
            repo_id=repo_id,
            repo_type="dataset",
            path_or_fileobj=BytesIO(
                task_results.model_dump_json(indent=2).encode("utf-8")
            ),
            path_in_repo=path_in_repo,
        )
        click.echo(f"Uploaded scores to hf://{repo_id}/{path_in_repo}")

        path_in_repo = f"{SUMMARIES_PREFIX}/{submission_path}/{SUMMARY_FILENAME}"
        hf_api.upload_file(
            repo_id=repo_id,
            repo_type="dataset",
            path_or_fileobj=BytesIO(stats.model_dump_json(indent=2).encode("utf-8")),
            path_in_repo=path_in_repo,
        )
        click.echo(f"Uploaded summary to hf://{repo_id}/{path_in_repo}")


cli.add_command(score_command)


@click.command(
    name="publish",
    help="Upload Inspect logs to HuggingFace for official scoring",
)
@click.argument("log_dir", type=click.Path(exists=True, file_okay=False))
@click.option(
    "--submissions-repo-id",
    type=str,
    default=lambda: os.environ.get("SUBMISSIONS_REPO_ID", ""),
    help="HF repo id for submissions. Defaults to SUBMISSIONS_REPO_ID env var.",
)
@click.option(
    "-o",
    "--openness",
    type=AliasedChoice(OPENNESS_MAPPING),
    required=True,
    help=generate_choice_help(OPENNESS_MAPPING, "Level of openness for the agent."),
)
@click.option(
    "-t",
    "--tool-usage",
    type=AliasedChoice(TOOL_MAPPING),
    required=True,
    help=generate_choice_help(TOOL_MAPPING, "Tool choices available to the agent."),
)
@click.option(
    "--username",
    type=str,
    default=None,
    help="HF username/org for submission. Defaults to your HF account name.",
)
@click.option(
    "--agent-name",
    type=str,
    required=True,
    help="Descriptive agent name for submission.",
)
@click.option(
    "--agent-description",
    type=str,
    default=None,
    help="Description of the agent being submitted.",
)
@click.option(
    "--agent-url",
    type=str,
    default=None,
    help="URL to the agent's repository or documentation.",
)
def publish_logs_command(
    log_dir: str,
    submissions_repo_id: str,
    openness: str,
    tool_usage: str,
    username: str | None,
    agent_name: str,
    agent_description: str | None,
    agent_url: str | None,
):
    # Allow huggingface imports to be optional
    from huggingface_hub import HfApi

    # Derive a filesafe agent_name
    safe_agent_name = sanitize_path_component(agent_name)
    if safe_agent_name != agent_name:
        click.echo(
            f"Note: agent_name '{agent_name}' contains unsafe characters; "
            f"using '{safe_agent_name}' for submission filenames."
        )

    with open(os.path.join(log_dir, EVAL_CONFIG_FILENAME), "r", encoding="utf-8") as f:
        eval_config = EvalConfig.model_validate_json(f.read())

    # Determine HF user
    hf_api = HfApi()
    if not username:
        try:
            username = hf_api.whoami()["name"]
            assert isinstance(username, str), "Invalid username type from HF API"
            click.echo(f"Defaulting username to Hugging Face account: {username}")
        except Exception:
            raise click.ClickException(
                "--username must be provided or ensure HF authentication is configured"
            )

    # Derive a filesafe username
    safe_username = sanitize_path_component(username)
    if safe_username != username:
        click.echo(
            f"Note: username '{username}' contains unsafe characters; "
            f"using '{safe_username}' for submission filenames."
        )

    # Fill submission metadata
    submission = SubmissionMetadata(
        username=username,
        agent_name=agent_name,
        agent_description=agent_description,
        agent_url=agent_url,
        submit_time=datetime.now(timezone.utc),
        openness=openness,
        tool_usage=tool_usage,
    )

    atomic_write_file(
        os.path.join(log_dir, SUBMISSION_METADATA_FILENAME),
        submission.model_dump_json(indent=2),
    )

    # Validate suite config version
    config_name = eval_config.suite_config.version
    if not config_name:
        raise click.ClickException("Suite config version is required for upload.")

    # Build submission name
    if submission.submit_time is None:
        raise click.ClickException("Submission timestamp is required for upload.")
    ts = submission.submit_time.strftime("%Y-%m-%dT%H-%M-%S")
    submission_name = f"{safe_username}_{safe_agent_name}_{ts}"

    # Upload logs and summary
    logs_url = upload_folder_to_hf(
        hf_api,
        log_dir,
        submissions_repo_id,
        config_name,
        eval_config.split,
        submission_name,
    )
    click.echo(f"Uploaded submission logs dir to {logs_url}")


cli.add_command(publish_logs_command)


@click.command(
    name="backfill",
    help="Backfill eval_config, scores, and submission files from legacy agenteval.json file",
)
@click.option(
    "--results-repo-id",
    type=str,
    required=False,
    default="allenai/asta-bench-internal-results",
)
@click.option(
    "--submissions-repo-id",
    type=str,
    required=False,
    default="allenai/asta-bench-internal-submissions",
)
@click.argument("submission_path", type=str)
def backfill_command(results_repo_id, submissions_repo_id, submission_path):
    with tempfile.TemporaryDirectory() as temp_dir:
        submissions_dir = os.path.join(temp_dir, "submissions")
        results_dir = os.path.join(temp_dir, "results")

        import huggingface_hub

        api = huggingface_hub.HfApi()
        api.snapshot_download(
            repo_id=results_repo_id,
            repo_type="dataset",
            allow_patterns=[f"{submission_path}.json"],
            local_dir=results_dir,
        )
        api.snapshot_download(
            repo_id=submissions_repo_id,
            repo_type="dataset",
            allow_patterns=[f"{submission_path}/agenteval.json"],
            local_dir=submissions_dir,
        )

        lb_submission = LeaderboardSubmission.model_validate_json(
            open(os.path.join(results_dir, f"{submission_path}.json")).read()
        )
        with open(
            os.path.join(submissions_dir, f"{submission_path}/submission.json"),
            "w",
            encoding="utf-8",
        ) as f:
            f.write(lb_submission.submission.model_dump_json(indent=2))
        with open(
            os.path.join(submissions_dir, f"{submission_path}/eval_config.json"),
            "w",
            encoding="utf-8",
        ) as f:
            eval_config = EvalConfig(
                suite_config=lb_submission.suite_config, split=lb_submission.split
            )
            f.write(eval_config.model_dump_json(indent=2))
        os.makedirs(
            os.path.join(submissions_dir, f"{SUMMARIES_PREFIX}/{submission_path}"),
            exist_ok=True,
        )
        with open(
            os.path.join(
                submissions_dir, f"{SUMMARIES_PREFIX}/{submission_path}/scores.json"
            ),
            "w",
            encoding="utf-8",
        ) as f:
            results = TaskResults(results=lb_submission.results)
            f.write(results.model_dump_json(indent=2))
        api.upload_folder(
            repo_id=submissions_repo_id,
            repo_type="dataset",
            folder_path=submissions_dir,
            path_in_repo="",
        )
        click.echo(f"Backfilled submission {submission_path} in {submissions_repo_id}")


cli.add_command(backfill_command)


@click.command(
    name="publish",
    help="Publish scored results in log_dir to HuggingFace leaderboard.",
)
@click.argument("submission_urls", nargs=-1, required=True, type=str)
@click.option(
    "--repo-id",
    default="allenai/asta-bench-internal-results",
    required=False,
    help="HuggingFace repo",
)
def publish_lb_command(repo_id: str, submission_urls: tuple[str, ...]):
    if not submission_urls:
        click.echo("At least one submission URL is required.")
        sys.exit(1)

    with tempfile.TemporaryDirectory() as temp_dir:
        from huggingface_hub import HfApi, snapshot_download

        local_submissions_dir = os.path.join(temp_dir, "submissions")
        local_results_dir = os.path.join(temp_dir, "results")

        hf_api = HfApi()

        try:
            paths_of_interest = RepoPathsOfInterest.from_urls(submission_urls)
        except ValueError as exc:
            click.echo(str(exc))
            sys.exit(1)

        submission_repo_id = paths_of_interest.repo_id
        submission_paths = paths_of_interest.relative_paths

        eval_config_rel_paths = [
            f"{p}/{EVAL_CONFIG_FILENAME}" for p in submission_paths
        ]
        scores_rel_paths = [
            f"{SUMMARIES_PREFIX}/{p}/{SCORES_FILENAME}" for p in submission_paths
        ]
        submission_metadata_rel_paths = [
            f"{p}/{SUBMISSION_METADATA_FILENAME}" for p in submission_paths
        ]

        # Download all input files in one shot
        snapshot_download(
            repo_id=submission_repo_id,
            repo_type="dataset",
            allow_patterns=eval_config_rel_paths
            + scores_rel_paths
            + submission_metadata_rel_paths,
            local_dir=local_submissions_dir,
        )

        lb_submissions = []
        # Create results files locally
        for (
            submission_url,
            submission_path,
            eval_config_path,
            scores_path,
            submission_metadata_path,
        ) in zip(
            submission_urls,
            submission_paths,
            eval_config_rel_paths,
            scores_rel_paths,
            submission_metadata_rel_paths,
        ):
            local_eval_config_path = os.path.join(
                local_submissions_dir, eval_config_path
            )
            local_scores_path = os.path.join(local_submissions_dir, scores_path)
            local_submission_path = os.path.join(
                local_submissions_dir, submission_metadata_path
            )
            required_files = [
                local_eval_config_path,
                local_scores_path,
                local_submission_path,
            ]

            missing = [
                os.path.basename(f) for f in required_files if not os.path.exists(f)
            ]
            if missing:
                click.echo(
                    "Skipping {}: missing {}".format(
                        submission_path, ", ".join(missing)
                    )
                )
                continue

            eval_config = EvalConfig.model_validate_json(
                open(local_eval_config_path).read()
            )
            results = TaskResults.model_validate_json(
                open(local_scores_path).read()
            ).results
            submission = SubmissionMetadata.model_validate_json(
                open(local_submission_path).read()
            )
            if not submission_url.startswith("hf://datasets/"):
                submission_url_to_use = submission_url.replace(
                    "hf://", "hf://datasets/", 1
                )
            else:
                submission_url_to_use = submission_url
            submission.logs_url = submission_url_to_use
            lb_submission = LeaderboardSubmission(
                suite_config=eval_config.suite_config,
                split=eval_config.split,
                results=results,
                submission=submission,
            )
            lb_submission = compress_model_usages(lb_submission)
            lb_submissions.append(lb_submission)
            os.makedirs(
                os.path.join(local_results_dir, os.path.dirname(submission_path)),
                exist_ok=True,
            )
            with open(
                os.path.join(local_results_dir, f"{submission_path}.json"),
                "w",
                encoding="utf-8",
            ) as f:
                f.write(lb_submission.model_dump_json(indent=None))

        # Validate the config with the schema in HF
        try:
            check_submissions_against_readme(
                lb_submissions=lb_submissions, repo_id=repo_id
            )
        except ValueError as exc:
            click.echo(str(exc))
            sys.exit(1)

        # Upload all results files in one shot
        click.echo(f"Uploading {len(submission_paths)} results to {repo_id}...")
        hf_api.upload_folder(
            folder_path=local_results_dir,
            path_in_repo="",
            repo_id=repo_id,
            repo_type="dataset",
        )
        click.echo("Done")


@click.group(name="lb", help="Leaderboard related commands")
def lb():
    pass


def validate_config(ctx, param, value):
    if value is not None:
        return value
    repo_id = ctx.params.get("repo_id")
    configs = datasets.get_dataset_config_names(repo_id)
    click.echo(f"Available configs: {configs}")
    click.echo("Please specify a config via --config")
    ctx.exit()


def validate_split(ctx, param, value):
    if value is not None:
        return value
    repo_id = ctx.params.get("repo_id")
    config = ctx.params.get("config")
    splits = datasets.get_dataset_split_names(repo_id, config_name=config)
    click.echo(f"Available splits: {splits}")
    click.echo("Please specify a split via --split")
    ctx.exit()


@lb.command(name="view", help="View leaderboard results.")
@click.option(
    "--repo-id",
    envvar="RESULTS_REPO_ID",
    required=True,
    help="HuggingFace dataset ID",
)
@click.option(
    "--config",
    default=None,
    callback=validate_config,
    help="Name of the dataset configuration to load",
)
@click.option(
    "--split",
    default=None,
    callback=validate_split,
    help="Dataset split to load",
)
@click.option(
    "--tag",
    default=None,
    help="If provided, show detail for this tag instead of overview",
)
@click.option(
    "--save-dir",
    default=None,
    type=click.Path(),
    help="Directory for saving plots (plots saved when specified)",
)
@click.option(
    "--save-no-subdirs",
    is_flag=True,
    default=False,
    help="Use exact save directory without auto-generated subdirectories",
)
@click.option(
    "--scatter-show-missing-cost",
    is_flag=True,
    default=False,
    help=(
        "Show agents with scores but no cost data in a separate section "
        "to the right of the dividing line in scatter plots"
    ),
)
@click.option(
    "--preserve-none-scores",
    is_flag=True,
    default=False,
    help="Preserve None values instead of treating them as 0 for incomplete scores",
)
@click.option(
    "--exclude-primary-metric",
    is_flag=True,
    default=False,
    help="Exclude primary metric (overall score+cost for overview, tag score+cost for tag views) from tables and plots",
)
@click.option(
    "--dedup",
    type=click.Choice(["index", "latest"]),
    default="index",
    help="How to handle duplicate agent names: 'index' (add numbers, default) or 'latest' (keep only latest)",
)
@click.option(
    "--exclude-agent-pattern",
    multiple=True,
    help="Regex pattern to exclude agents by name/model (case-insensitive, can be specified multiple times)",
)
@click.option(
    "--include-task-pattern",
    multiple=True,
    help="Regex pattern to include only matching tasks/sub-benchmarks by name (case-insensitive, can be specified multiple times). Only applies when --tag is specified.",
)
@click.option(
    "--scatter-legend-max-width",
    type=int,
    default=None,
    help="Maximum width in characters for scatter legend text wrapping (default: no wrapping). Set to enable wrapping, e.g., 35.",
)
@click.option(
    "--scatter-figure-width",
    type=float,
    default=None,
    help="Total scatter figure width in inches including legend space (default: auto-calculate based on subplots).",
)
@click.option(
    "--scatter-subplot-height",
    type=float,
    default=None,
    help="Height of each scatter subplot in inches (default: matplotlib default aspect ratio).",
)
@click.option(
    "--scatter-subplot-spacing",
    type=float,
    default=None,
    show_default=True,
    help="Vertical spacing between scatter subplots as fraction of subplot height (default: matplotlib default).",
)
@click.option(
    "--scatter-x-log-scale",
    is_flag=True,
    help="Use log scale for x-axis in scatter plots.",
)
@click.option(
    "--is-internal",
    is_flag=True,
    default=False,
    help="Show internal logs URLs instead of public ones.",
)
def view_command(
    repo_id,
    config,
    split,
    tag,
    save_dir,
    save_no_subdirs,
    scatter_show_missing_cost,
    preserve_none_scores,
    exclude_primary_metric,
    dedup,
    exclude_agent_pattern,
    include_task_pattern,
    scatter_legend_max_width,
    scatter_figure_width,
    scatter_subplot_height,
    scatter_subplot_spacing,
    scatter_x_log_scale,
    is_internal,
):
    """View a specific config and split; show overview or tag detail."""
    from .leaderboard.view import LeaderboardViewer

    viewer = LeaderboardViewer(repo_id, config, split, is_internal=is_internal)

    df, plots = viewer.view(
        tag,
        with_plots=bool(save_dir),
        preserve_none_scores=preserve_none_scores,
        exclude_primary_metric=exclude_primary_metric,
        duplicate_handling=dedup,
        exclude_agent_patterns=(
            list(exclude_agent_pattern) if exclude_agent_pattern else None
        ),
        include_task_patterns=(
            list(include_task_pattern) if include_task_pattern else None
        ),
        scatter_show_missing_cost=scatter_show_missing_cost,
        scatter_legend_max_width=scatter_legend_max_width,
        scatter_figure_width=scatter_figure_width,
        scatter_subplot_height=scatter_subplot_height,
        scatter_subplot_spacing=scatter_subplot_spacing,
        scatter_x_log_scale=scatter_x_log_scale,
    )
    click.echo(df.to_string(index=False))

    if save_dir:
        if not save_no_subdirs:
            subdir = tag or "overall"
            ts = datetime.now().strftime("%Y%m%d_%H%M%S")
            safe_repo = repo_id.replace("/", "_")
            base = save_dir
            sub = f"{safe_repo}_{config}_{split}"
            outdir = os.path.join(base, sub, f"{subdir}_{ts}")
        else:
            outdir = save_dir
        os.makedirs(outdir, exist_ok=True)

        jsonl_path = os.path.join(outdir, "data.jsonl")
        df.to_json(jsonl_path, orient="records", lines=True)
        click.echo(f"Saved data: {jsonl_path}")

        for name, fig in plots.items():
            path = os.path.join(outdir, f"{name}.png")
            if "scatter" in name and scatter_figure_width is not None:
                # When custom width is specified (likely for publication), use higher DPI
                fig.savefig(path, bbox_inches="tight", dpi=300)
            else:
                fig.savefig(path, bbox_inches="tight")
            click.echo(f"Saved plot: {path}")


lb.add_command(publish_lb_command)
cli.add_command(lb)


@cli.command(
    name="eval",
    help="Run inspect eval-set on specified tasks with the given arguments",
    context_settings={"ignore_unknown_options": True},
)
@click.option(
    "--log-dir",
    type=str,
    help="Log directory. Defaults to INSPECT_LOG_DIR or auto-generated under ./logs.",
)
@click.option(
    "--config-path",
    "config_path",
    type=str,
    help="Path to a yml config file.",
    required=True,
)
@click.option(
    "--split",
    type=str,
    help="Config data split.",
    required=True,
)
@click.option(
    "--ignore-git",
    is_flag=True,
    help="Ignore git reproducibility checks (not recommended).",
)
@click.option(
    "--config-only",
    is_flag=True,
    help="Print the command that would be run and save eval_config locally.",
)
@click.option(
    "--display",
    type=str,
    # https://github.com/UKGovernmentBEIS/inspect_ai/issues/1891 and
    # https://github.com/allenai/nora-issues-research/issues/77#issuecomment-2877262319
    # TODO: remove this once fixed
    help="Display format. Defaults to plain.",
    default="plain",
)
@click.argument("args", nargs=-1, type=click.UNPROCESSED)
def eval_command(
    log_dir: str | None,
    config_path: str,
    split: str,
    ignore_git: bool,
    config_only: bool,
    display: str,
    args: tuple[str],
):
    """Run inspect eval-set with arguments and append tasks"""
    suite_config = load_suite_config(config_path)
    tasks = suite_config.get_tasks(split)

    # Verify git status for reproducibility
    if not ignore_git:
        verify_git_reproducibility()

    if not log_dir:
        log_dir = os.environ.get("INSPECT_LOG_DIR")
        if not log_dir:
            timestamp = datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
            log_dir = os.path.join(
                ".",
                "logs",
                f"{suite_config.name}_{suite_config.version}_{split}_{timestamp}",
            )
            click.echo(f"No log dir was manually set; using {log_dir}")
    logd_args = ["--log-dir", log_dir]
    display_args = ["--display", display]

    # Write the config portion of the results file
    os.makedirs(log_dir, exist_ok=True)

    inspect_command = (
        ["inspect", "eval-set"]
        + list(args)
        + logd_args
        + display_args
        + [x.path for x in tasks]
    )

    eval_config = EvalConfig(
        suite_config=suite_config, split=split, inspect_command=inspect_command
    )

    eval_config_path = os.path.join(log_dir, EVAL_CONFIG_FILENAME)
    if not os.path.exists(eval_config_path):
        with open(eval_config_path, "w", encoding="utf-8") as f:
            f.write(eval_config.model_dump_json(indent=2))
    else:
        with open(eval_config_path, "r", encoding="utf-8") as f:
            existing_config = EvalConfig.model_validate_json(f.read())
        if existing_config != eval_config:
            click.echo(
                f"Suite config does not match pre-existing config in {EVAL_CONFIG_FILENAME}. Rerun in an empty directory"
            )
            sys.exit(1)

    # We use subprocess here to keep arg management simple; an alternative
    # would be calling `inspect_ai.eval_set()` directly, which would allow for
    # programmatic execution
    if config_only:
        click.echo(f"Dry run: would run command: {' '.join(inspect_command)}")
        return

    click.echo(f"Running {' '.join(inspect_command)}")
    proc = subprocess.run(inspect_command)

    if proc.returncode != 0:
        raise click.ClickException(
            f"inspect eval-set failed while running {config_path}"
        )

    ctx = click.get_current_context()
    click.echo(
        f"You can now run '{ctx.parent.info_name if ctx.parent else 'cli'} score {log_dir}' to score the results"
    )


cli.add_command(eval_command)


@dataclass
class WithinRepoPathComponents:
    hf_config: str
    split: str
    submission_name: str

    def to_within_repo_result_path(self):
        return f"{self.hf_config}/{self.split}/{self.submission_name}.json"

    def to_within_repo_submission_dir(self):
        return f"{self.hf_config}/{self.split}/{self.submission_name}"

    def to_within_repo_submission_summary_dir(self):
        return f"summaries/{self.hf_config}/{self.split}/{self.submission_name}"

    def to_within_repo_submission_eval_config(self):
        return f"{self.to_within_repo_submission_dir()}/{EVAL_CONFIG_FILENAME}"

    def to_within_repo_submission_submission_metadata(self):
        return f"{self.to_within_repo_submission_dir()}/{SUBMISSION_METADATA_FILENAME}"

    def to_within_repo_submission_scores(self):
        return f"{self.to_within_repo_submission_summary_dir()}/{SCORES_FILENAME}"

    def within_repo_submission_patterns(self):
        return [
            self.to_within_repo_submission_eval_config(),
            self.to_within_repo_submission_submission_metadata(),
            self.to_within_repo_submission_scores(),
            f"{self.to_within_repo_submission_dir()}/*.eval",
        ]

    @staticmethod
    def from_within_repo_result_path(result_path):
        [hf_config, split, filename] = result_path.split("/")
        suffix = ".json"
        assert filename.endswith(suffix)
        submission_name = filename[: -len(suffix)]
        return WithinRepoPathComponents(
            hf_config=hf_config, split=split, submission_name=submission_name
        )

    @staticmethod
    def from_within_repo_submission_path(submission_path):
        [hf_config, split, submission_name] = submission_path.split("/")
        return WithinRepoPathComponents(
            hf_config=hf_config, split=split, submission_name=submission_name
        )


@click.command(name="copy", help="Copy a result from one results HF repo to another.")
@click.argument("result_urls", nargs=-1, required=True, type=str)
@click.option(
    "--target-submissions-repo",
    default=None,
    required=False,
    help="Provide this if you also want to copy the underlying submission to another submissions HF repo.",
)
@click.option(
    "--target-results-repo",
    default=None,
    required=True,
)
@click.option(
    "--read-public-logs-field",
    is_flag=True,
    default=False,
    help="Provide this if the source results have log urls in logs_url_public.",
)
@click.option(
    "--write-public-logs-field",
    is_flag=True,
    default=False,
    help="Provide this if the target results should have log urls in logs_url_public.",
)
def copy_command(
    target_submissions_repo: str | None,
    target_results_repo: str,
    read_public_logs_field: bool,
    write_public_logs_field: bool,
    result_urls: tuple[str, ...],
):
    try:
        src_result_paths_of_interest = RepoPathsOfInterest.from_urls(result_urls)
    except ValueError as exc:
        click.echo(str(exc))
        sys.exit(1)

    src_results_repo = src_result_paths_of_interest.repo_id
    result_paths = src_result_paths_of_interest.relative_paths
    click.echo(f"{len(result_paths)} result files to copy.")

    click.echo(f"source results repo: {src_results_repo}")
    click.echo(f"target results repo: {target_results_repo}")

    with tempfile.TemporaryDirectory() as temp_dir:
        from huggingface_hub import HfApi, snapshot_download

        hf_api = HfApi()

        local_src_results_dir = os.path.join(temp_dir, "sourceresults")
        local_target_results_dir = os.path.join(temp_dir, "targetresults")

        snapshot_download(
            repo_id=src_results_repo,
            repo_type="dataset",
            allow_patterns=result_paths,
            local_dir=local_src_results_dir,
        )

        log_urls = []
        lb_submissions = []

        for result_path in result_paths:
            local_src_result_path = os.path.join(local_src_results_dir, result_path)
            with open(local_src_result_path) as f_src_result:
                result = LeaderboardSubmission.model_validate(json.load(f_src_result))

            current_logs_url = (
                result.submission.logs_url_public
                if read_public_logs_field
                else result.submission.logs_url
            )
            log_urls.append(current_logs_url)
            _, src_submission_path = parse_hf_url(current_logs_url)

            if target_submissions_repo is not None:
                new_logs_url = (
                    f"hf://datasets/{target_submissions_repo}/{src_submission_path}"
                )
                if write_public_logs_field:
                    result.submission.logs_url_public = new_logs_url
                    result.submission.logs_url = None
                else:
                    result.submission.logs_url_public = None
                    result.submission.logs_url = new_logs_url

            lb_submissions.append(result)

            os.makedirs(
                os.path.join(local_target_results_dir, os.path.dirname(result_path)),
                exist_ok=True,
            )
            with open(
                os.path.join(local_target_results_dir, result_path),
                "w",
                encoding="utf-8",
            ) as f_target_result:
                f_target_result.write(result.model_dump_json(indent=None))

        # Validate the config with the schema in HF
        check_submissions_against_readme(
            lb_submissions=lb_submissions, repo_id=target_results_repo
        )
        click.echo(
            f"Uploading {len(lb_submissions)} results to {target_results_repo}..."
        )
        hf_api.upload_folder(
            folder_path=local_target_results_dir,
            path_in_repo="",
            repo_id=target_results_repo,
            repo_type="dataset",
        )

        if target_submissions_repo is not None:
            try:
                src_submission_paths_of_interest = RepoPathsOfInterest.from_urls(
                    log_urls
                )
            except ValueError as exc:
                click.echo(str(exc))
                sys.exit(1)

            src_submissions_repo = src_submission_paths_of_interest.repo_id
            submission_paths = src_submission_paths_of_interest.relative_paths

            # I think we can just use the same local dir.
            local_submissions_dir = os.path.join(temp_dir, "submissions")
            paths_to_pull = []
            for submission_path in submission_paths:
                paths_to_pull.extend(
                    WithinRepoPathComponents.from_within_repo_submission_path(
                        submission_path
                    ).within_repo_submission_patterns()
                )

            snapshot_download(
                repo_id=src_submissions_repo,
                repo_type="dataset",
                allow_patterns=paths_to_pull,
                local_dir=local_submissions_dir,
            )
            click.echo(
                f"Uploading {len(submission_paths)} submissions to {target_submissions_repo}..."
            )
            hf_api.upload_folder(
                folder_path=local_submissions_dir,
                path_in_repo="",
                repo_id=target_submissions_repo,
                repo_type="dataset",
            )

    click.echo("Done")


cli.add_command(copy_command)


if __name__ == "__main__":
    cli()