"""Workspace layout: iteration directories and per-case output writers.

Outputs land under ``<workspace>/iteration-N/eval-<id>-<target>/<config>/`` with
``outputs/response.md``, ``timing.json``, and ``grading.json`` — the layout the
benchmark aggregator reads back.
"""

import json
from pathlib import Path
from typing import Any

from langchain_core.language_models import BaseChatModel

from skill_eval_runner.agent import AgentResult, run_agent
from skill_eval_runner.config import EvalSuiteConfig
from skill_eval_runner.grading import GradingResult, grade
from skill_eval_runner.pricing import calculate_cost
from skill_eval_runner.suite import EvalCase


def resolve_iteration(workspace_dir: Path, force: int | None) -> Path:
    """Return the iteration directory, creating it (auto-incrementing unless forced)."""
    workspace_dir.mkdir(parents=True, exist_ok=True)
    if force is not None:
        directory = workspace_dir / f"iteration-{force}"
        directory.mkdir(exist_ok=True)
        return directory
    existing = sorted(
        workspace_dir.glob("iteration-*"), key=lambda p: int(p.name.split("-")[1])
    )
    number = (int(existing[-1].name.split("-")[1]) + 1) if existing else 1
    directory = workspace_dir / f"iteration-{number}"
    directory.mkdir()
    return directory


def run_and_save(
    case: EvalCase,
    config_name: str,
    iteration_dir: Path,
    model: BaseChatModel,
    model_name: str,
    grader_model: BaseChatModel,
    grader_model_name: str,
    config: EvalSuiteConfig,
) -> float | None:
    """Run one case in one config, persist outputs + grading, and return its cost."""
    with_skill = config_name == "with_skill"
    out_dir = iteration_dir / case.dir_name() / config_name
    (out_dir / "outputs").mkdir(parents=True, exist_ok=True)

    print(f"  [{config_name}] running agent...", flush=True)
    result = run_agent(case, with_skill=with_skill, model=model, config=config)
    total_tokens = result.input_tokens + result.output_tokens
    cost_usd = calculate_cost(
        model_name, result.input_tokens, result.output_tokens, config.pricing
    )

    cost_str = f"${cost_usd:.4f}" if cost_usd is not None else "$unknown"
    print(
        f"  [{config_name}] done — {result.duration_ms / 1000:.1f}s, "
        f"{total_tokens:,} tokens ({cost_str})",
        flush=True,
    )
    (out_dir / "outputs" / config.artifact.filename).write_text(result.response)

    print(f"  [{config_name}] grading...", flush=True)
    graded = grade(result.response, case.assertions, grader_model)
    (out_dir / "grading.json").write_text(json.dumps(graded.grading, indent=2))
    _print_score(config_name, graded.grading)

    # The grader is a second model call (possibly a different model); price it on its
    # own model and report the agent + grader sum as the case's total cost.
    grader_cost = calculate_cost(
        grader_model_name, graded.input_tokens, graded.output_tokens, config.pricing
    )
    total_cost = _combine_costs(cost_usd, grader_cost)
    _print_grader(config_name, graded, grader_cost)

    (out_dir / "timing.json").write_text(
        json.dumps(
            _timing(
                model_name,
                grader_model_name,
                result,
                cost_usd,
                graded,
                grader_cost,
                total_cost,
            ),
            indent=2,
        )
    )
    return total_cost


def grade_only(
    case: EvalCase,
    config_name: str,
    iteration_dir: Path,
    grader_model: BaseChatModel,
    grader_model_name: str,
    config: EvalSuiteConfig,
) -> float | None:
    """Re-grade an existing response without re-running the agent; return grader cost."""
    out_dir = iteration_dir / case.dir_name() / config_name
    response_path = out_dir / "outputs" / config.artifact.filename
    if not response_path.exists():
        print(f"  [{config_name}] no output found, skipping")
        return None
    graded = grade(response_path.read_text(), case.assertions, grader_model)
    (out_dir / "grading.json").write_text(json.dumps(graded.grading, indent=2))
    _print_score(config_name, graded.grading)

    grader_cost = calculate_cost(
        grader_model_name, graded.input_tokens, graded.output_tokens, config.pricing
    )
    _print_grader(config_name, graded, grader_cost)
    _update_timing_grader(
        out_dir / "timing.json", grader_model_name, graded, grader_cost
    )
    return grader_cost


def _combine_costs(*parts: float | None) -> float | None:
    """Sum the known costs; ``None`` only when every part is unpriced (None)."""
    known = [p for p in parts if p is not None]
    return round(sum(known), 6) if known else None


def _timing(
    model_name: str,
    grader_model_name: str,
    result: AgentResult,
    cost_usd: float | None,
    graded: GradingResult,
    grader_cost: float | None,
    total_cost: float | None,
) -> dict[str, Any]:
    """Assemble the per-case ``timing.json`` payload (agent + grader, split + summed)."""
    return {
        "model": model_name,
        "input_tokens": result.input_tokens,
        "output_tokens": result.output_tokens,
        "total_tokens": result.input_tokens + result.output_tokens,
        "duration_ms": result.duration_ms,
        "cost_usd": cost_usd,
        "grader_model": grader_model_name,
        "grader_input_tokens": graded.input_tokens,
        "grader_output_tokens": graded.output_tokens,
        "grader_total_tokens": graded.input_tokens + graded.output_tokens,
        "grader_cost_usd": grader_cost,
        "total_cost_usd": total_cost,
    }


def _update_timing_grader(
    timing_path: Path,
    grader_model_name: str,
    graded: GradingResult,
    grader_cost: float | None,
) -> None:
    """Merge grader usage/cost into an existing ``timing.json`` (re-grade path)."""
    if not timing_path.exists():
        return
    timing = json.loads(timing_path.read_text())
    timing["grader_model"] = grader_model_name
    timing["grader_input_tokens"] = graded.input_tokens
    timing["grader_output_tokens"] = graded.output_tokens
    timing["grader_total_tokens"] = graded.input_tokens + graded.output_tokens
    timing["grader_cost_usd"] = grader_cost
    timing["total_cost_usd"] = _combine_costs(timing.get("cost_usd"), grader_cost)
    timing_path.write_text(json.dumps(timing, indent=2))


def _print_score(config_name: str, grading: dict[str, object]) -> None:
    """Print the ``passed/total (rate%)`` line for one config's grading."""
    summary = grading["summary"]
    assert isinstance(summary, dict)
    print(
        f"  [{config_name}] {summary['passed']}/{summary['total']} "
        f"assertions passed ({summary['pass_rate']:.0%})",
        flush=True,
    )


def _print_grader(
    config_name: str, graded: GradingResult, grader_cost: float | None
) -> None:
    """Print the grader's own token/cost line so its overhead is visible per case."""
    grader_total = graded.input_tokens + graded.output_tokens
    cost_str = f"${grader_cost:.4f}" if grader_cost is not None else "$unknown"
    print(
        f"  [{config_name}] grader: {grader_total:,} tokens ({cost_str})",
        flush=True,
    )