"""abreka test — agent validates experiment code and runs tests."""

from __future__ import annotations

import time
from datetime import datetime, timezone

import click
from rich.console import Console

from abreka.config import find_abreka_dir, load_config
from abreka.display import stream_event
from abreka.models import ExperimentStatus
from abreka.persistence import (
    experiment_dir,
    load_experiment,
    normalize_exp_id,
    save_experiment,
    update_index_entry,
)
from abreka.pi_client import PiRpcClient
from abreka.prompts import (
    format_wiki_index,
    load_template,
    render_template,
)
from abreka.wiki import load_wiki_index
from abreka.stage_sessions import stage_resume_message, stage_session_path
from abreka.runtime_costs import record_stage_runtime


@click.command("test")
@click.argument("exp_id")
@click.option("--resume", is_flag=True, help="Resume the existing test session.")
def test_cmd(exp_id: str, resume: bool):
    """Agent validates experiment code and runs the test suite."""
    abreka_dir = find_abreka_dir()
    config = load_config(abreka_dir)
    exp_id = normalize_exp_id(exp_id)
    exp = load_experiment(abreka_dir, exp_id)
    session_path = stage_session_path(abreka_dir, exp_id, "test")

    if resume:
        if exp.status != ExperimentStatus.TESTING:
            raise click.ClickException(
                f"Experiment {exp_id} is '{exp.status.value}', expected 'testing'."
            )
        if not session_path.is_file():
            raise click.ClickException(
                f"No existing test session found for experiment {exp_id}."
            )
    else:
        if exp.status not in (ExperimentStatus.IMPLEMENTING, ExperimentStatus.TESTING):
            raise click.ClickException(
                f"Experiment {exp_id} is '{exp.status.value}', expected 'implementing' or 'testing'."
            )
        if session_path.exists():
            raise click.ClickException(
                f"Test session already exists for experiment {exp_id}. Use --resume to continue it."
            )

    console = Console()
    console.print(f"[bold]Testing experiment {exp_id}...[/bold]\n")

    started_at = datetime.now(timezone.utc)
    started_monotonic = time.monotonic()
    stage_status = "ok"

    try:
        if not resume:
            exp.status = ExperimentStatus.TESTING
            save_experiment(exp, abreka_dir)
            update_index_entry(abreka_dir, exp)

        template = load_template(abreka_dir, "test")
        prompt_text = render_template(template, {
            "exp_id": exp_id,
            "hypothesis": exp.hypothesis,
            "method": exp.method,
            "lib_name": config.project.lib,
            "test_command": config.project.test_command,
            "wiki_index": format_wiki_index(load_wiki_index(abreka_dir)),
        })
        message = stage_resume_message("test", exp_id) if resume else prompt_text

        model = config.agent.model

        with PiRpcClient(
            provider=config.agent.provider,
            model=model,
            cwd=abreka_dir.parent,
            thinking_level=config.agent.thinking_level,
            extra_args=config.agent.pi_extra_args,
            session_path=session_path,
            session_capture_mode=config.agent.session_capture_mode,
        ) as pi:
            pi.prompt(message, on_event=stream_event)

        console.print(f"\n[bold green]Testing complete for {exp_id}[/bold green]")
    except Exception:
        if stage_status == "ok":
            stage_status = "error"
        raise
    finally:
        record_stage_runtime(
            abreka_dir,
            exp_id,
            stage="test",
            elapsed_seconds=time.monotonic() - started_monotonic,
            status=stage_status,
            started_at=started_at,
            finished_at=datetime.now(timezone.utc),
        )
