"""Helpers for native pi session files used by agent stages."""

from __future__ import annotations

import json
from pathlib import Path

from abreka.persistence import experiment_dir


STAGE_SESSION_FILENAMES: dict[str, str] = {
    "propose": "pi_session_propose.jsonl",
    "implement": "pi_session_implement.jsonl",
    "test": "pi_session_test.jsonl",
    "evaluate": "pi_session_evaluate.jsonl",
}


def stage_session_filename(stage: str) -> str:
    return STAGE_SESSION_FILENAMES.get(stage, f"pi_session_{stage}.jsonl")


def stage_session_path(abreka_dir: Path, exp_id: str, stage: str) -> Path:
    return experiment_dir(abreka_dir, exp_id) / "artifacts" / stage_session_filename(stage)


def stage_resume_message(stage: str, exp_id: str) -> str:
    if stage == "propose":
        return (
            f"Resume the interrupted propose step for experiment {exp_id}. "
            f"The reserved experiment already exists. Inspect the repo state and prior session work, "
            f"then finish filling out experiment {exp_id} and set its status to proposed. "
            "Do not create another experiment."
        )
    return (
        f"Resume the interrupted {stage} step for experiment {exp_id}. "
        "First inspect the current repo state and the prior session work so you do not duplicate completed work. "
        f"Then finish the remaining {stage} work."
    )


def load_session_entries(path: Path) -> list[dict]:
    if not path.is_file():
        return []
    entries: list[dict] = []
    with open(path) as f:
        for line in f:
            stripped = line.strip()
            if not stripped:
                continue
            entries.append(json.loads(stripped))
    return entries
