import asyncio
import argparse
import json
import os

from vals import configure_credentials
from vals.sdk.custom_metric import CustomMetric
from vals.sdk.run import Run


def extract_run_ids_from_runs_json(
    runs_json: dict,
    limit: int = -1,
    categories: list[str] | None = None,
) -> list[str]:
    """
    Extract run IDs from a runs JSON structure.

    Expected shape: { "<model_name>": { "<category>": ["<run_id>", ...], ... }, ... }

    Args:
        runs_json: The loaded runs JSON.
        limit: Max number of run IDs to return (-1 = no limit).
        categories: If set, only include run IDs from these category keys (e.g. ["stepwise_reasoning"]).
    """
    run_ids: list[str] = []
    seen: set[str] = set()
    include_categories = set(categories) if categories else None
    for model_data in runs_json.values():
        if not isinstance(model_data, dict):
            continue
        for category, run_list in model_data.items():
            if include_categories is not None and category not in include_categories:
                continue
            if not isinstance(run_list, list):
                continue
            for run_id in run_list:
                if isinstance(run_id, str):
                    rid = run_id.strip()
                    if rid and rid not in seen:
                        seen.add(rid)
                        run_ids.append(rid)
    return run_ids[:limit]


async def _run_metric_on_one(
    metric: CustomMetric,
    run_id: str,
    semaphore: asyncio.Semaphore,
) -> tuple[str, str | None, str | None]:
    """Run metric on a single run; returns (run_id, result, error). One of result/error is None."""
    async with semaphore:
        try:
            result = await metric.run(run_id=run_id, persist=True)
            return (run_id, result, None)
        except Exception as e:
            return (run_id, None, str(e))


async def add_metric_to_suite_runs(
    metric_id: str,
    *,
    suite_id: str | None = None,
    run_ids: list[str] | None = None,
    runs_json: dict | None = None,
    categories: list[str] | None = None,
    project_id: str = "default-project",
    limit: int | None = None,
    concurrency: int = 10,
):
    """
    Run an existing custom metric on runs and persist the results.

    Provide run_ids, runs_json (model->category->[run_ids] dict), or suite_id
    (with optional project_id) to run the metric on all runs in that suite.

    Args:
        metric_id: The ID of the custom metric to run
        suite_id: The ID of the test suite (required if run_ids/runs_json not provided)
        run_ids: Optional list of specific run IDs to add the metric to
        runs_json: Optional dict of runs (model name -> category -> list of run IDs); IDs are extracted and used
        categories: When using runs_json, only include run IDs from these category keys (e.g. ["stepwise_reasoning"])
        project_id: The project ID (defaults to "default-project", used when listing by suite)
        limit: Optional max number of runs to process (and, when using suite_id, max runs to fetch from suite)
        concurrency: Max number of runs to process in parallel (default 10).
    """
    # Extract run IDs from runs_json if provided
    if runs_json is not None:
        if categories:
            print(f"Filtering to categories: {categories}")
        run_ids = extract_run_ids_from_runs_json(runs_json, limit, categories)
    
    # Raise an error if no run IDs are provided
    if not run_ids and not suite_id:
        raise ValueError("Provide either run_ids, runs_json, or suite_id")

    # Get the custom metric
    metric = await CustomMetric.from_id(metric_id)
    print(f"Loaded metric: {metric.name}")

    if run_ids:
        runs_to_process = [run_id.strip() for run_id in run_ids]
        print(f"Targeting {len(runs_to_process)} run(s) by ID")
    else:
        # Get runs from the suite
        list_kwargs: dict = {"suite_id": suite_id, "project_id": project_id}
        if limit is not None:
            list_kwargs["limit"] = limit
        runs = await Run.list_runs(**list_kwargs)
        runs_to_process = [run_meta.id for run_meta in runs]
        print(f"Found {len(runs_to_process)} runs in suite")
        print([run.name for run in runs])

    if limit is not None:
        runs_to_process = runs_to_process[:limit]
        print(f"Limiting to {len(runs_to_process)} run(s)")

    total = len(runs_to_process)
    semaphore = asyncio.Semaphore(concurrency)
    print(f"Processing {total} run(s) with concurrency {concurrency}...")

    # Create tasks and print results as they complete so you still get
    # streaming output even when running concurrently.
    tasks = [
        asyncio.create_task(_run_metric_on_one(metric, run_id, semaphore))
        for run_id in runs_to_process
    ]

    succeeded = 0
    failed = 0
    completed = 0

    for finished in asyncio.as_completed(tasks):
        run_id, result, error = await finished
        completed += 1
        prefix = f"[{completed}/{total}] "
        if error is not None:
            failed += 1
            print(f"{prefix}{run_id}: Error: {error}", flush=True)
        else:
            succeeded += 1
            print(f"{prefix}{run_id}: {result}", flush=True)

    print(f"Done: {succeeded} succeeded, {failed} failed.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Run a custom metric on runs (all runs in a suite or specific run IDs)"
    )
    parser.add_argument("--metric-id", required=True, help="ID of the custom metric to run")
    parser.add_argument(
        "--suite-id",
        help="ID of the test suite (use this to run on all runs in the suite)",
    )
    parser.add_argument(
        "--run-ids",
        help="Comma-separated run IDs to add the metric to (alternative to --suite-id)",
    )
    parser.add_argument(
        "--runs-json",
        help="Path to JSON file of runs (model -> category -> list of run IDs); run IDs are extracted and used",
    )
    parser.add_argument(
        "--categories",
        help="When using --runs-json, only use run IDs from these categories (comma-separated, e.g. stepwise_reasoning)",
    )
    parser.add_argument(
        "--project-id",
        default="default-project",
        help="Project ID when using --suite-id (default: default-project)",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        metavar="N",
        help="Max number of runs to process (and to fetch when using --suite-id)",
    )
    parser.add_argument(
        "--concurrency",
        type=int,
        default=1,
        metavar="N",
        help="Max number of runs to process in parallel (default: 10)",
    )
    args = parser.parse_args()

    if not args.run_ids and not args.runs_json and not args.suite_id:
        parser.error("Provide one of: --suite-id, --run-ids, or --runs-json")

    run_ids: list[str] | None = None
    runs_json: dict | None = None
    categories: list[str] | None = None
    if args.run_ids:
        run_ids = [s.strip() for s in args.run_ids.split(",")]
    elif args.runs_json:
        with open(args.runs_json, encoding="utf-8") as f:
            runs_json = json.load(f)
    if args.categories:
        categories = [s.strip() for s in args.categories.split(",") if s.strip()]

    configure_credentials(api_key=os.environ["VALS_API_KEY"])

    asyncio.run(
        add_metric_to_suite_runs(
            args.metric_id,
            suite_id=args.suite_id,
            run_ids=run_ids,
            runs_json=runs_json,
            categories=categories,
            project_id=args.project_id,
            limit=args.limit,
            concurrency=args.concurrency,
        )
    )
