"""Save TAP run results (result.json only)."""

import json
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional


@dataclass
class RunResult:
    """Result of a TAP run. Returned by TAP.run()."""

    result_prompt: Optional[str]
    iteration_log: List[Dict[str, Any]]
    target_response: Optional[str]
    extra: Any  # None for general jailbreaking; custom judge may set e.g. a deal [A, B]
    success: bool


def _build_summary(
    success: bool,
    iteration_log: List[Dict[str, Any]],
    result_prompt: Optional[str],
) -> str:
    """One-line summary of why the run ended (for result.json)."""
    if success and result_prompt:
        return "Jailbreak found; winning prompt saved."
    if not iteration_log:
        return "Run stopped: no iterations (no branches)."
    last = iteration_log[-1]
    num_on = last.get("num_on_topic")
    num_off = last.get("num_off_topic")
    if num_on is not None and num_on == 0 and (num_off or 0) > 0:
        n = last.get("iteration", "?")
        return f"Stopped after iteration {n}: all branches pruned (off-topic). No target queries."
    depth = len(iteration_log)
    return f"Stopped: max depth ({depth}) reached; no jailbreak. Check iteration_log for scores and candidates."


def get_results_dir() -> Path:
    """Results directory: ./results/ in the current working directory."""
    return Path.cwd() / "results"


def get_single_dir() -> Path:
    """Directory for single-run results: results/single/."""
    return get_results_dir() / "single"


def get_batch_dir() -> Path:
    """Directory for batch results: results/batch/."""
    return get_results_dir() / "batch"


def make_run_id() -> str:
    """Unique run ID based on timestamp."""
    return datetime.now().strftime("%Y%m%d_%H%M%S")


def save_result(
    run_id: str,
    goal: str,
    config_name: str,
    config: Dict[str, Any],
    run_result: RunResult,
    run_dir: Optional[Path] = None,
) -> Path:
    """
    Write result.json for this run. Pass the RunResult returned by TAP.run().
    Returns the run directory path.
    """
    if run_dir is None:
        run_dir = get_single_dir() / run_id
    run_dir.mkdir(parents=True, exist_ok=True)

    summary = _build_summary(
        success=run_result.success,
        iteration_log=run_result.iteration_log,
        result_prompt=run_result.result_prompt,
    )

    payload: Dict[str, Any] = {
        "run_id": run_id,
        "goal": goal,
        "config_name": config_name,
        "config": config,
        "success": run_result.success,
        "summary": summary,
        "result_prompt": run_result.result_prompt,
        "target_response": run_result.target_response,
        "result_deal": run_result.extra,
        "iteration_log": run_result.iteration_log,
    }
    with open(run_dir / "result.json", "w") as f:
        json.dump(payload, f, indent=2)

    return run_dir
