from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import pandas as pd  # type: ignore[import-untyped]

from ..persistence.repository import Repository


def get_repo(db_path: str | Path) -> Repository:
    return Repository(db_path=db_path)


def list_runs_df(repo: Repository) -> pd.DataFrame:
    return pd.DataFrame(repo.list_runs())


def get_run_overview(repo: Repository, run_id: str) -> dict[str, Any]:
    run = repo.get_run(run_id)
    return {
        "name": run["name"],
        "started_at": run["started_at"],
        "completed_at": run["completed_at"],
        "status": run["status"],
        "total_cost_usd": run["total_cost_usd"],
        "config": json.loads(run["config_json"]),
    }


def generations_df(repo: Repository, run_id: str) -> pd.DataFrame:
    return pd.DataFrame(repo.list_generations(run_id))


def evaluations_df(repo: Repository, run_id: str) -> pd.DataFrame:
    return pd.DataFrame(repo.list_evaluations(run_id))


def genomes_df(
    repo: Repository, run_id: str, generation_idx: int | None = None
) -> pd.DataFrame:
    rows = repo.list_genomes(run_id, generation_idx)
    flat: list[dict[str, Any]] = []
    for r in rows:
        payload = json.loads(r["payload_json"])
        flat.append(
            {
                "id": r["id"],
                "generation_idx": r["generation_idx"],
                **payload,
            }
        )
    return pd.DataFrame(flat)