#!/usr/bin/env python3
"""Project-local Agent tuning test runner.

Generated by the atk-init Skill.
- Preserves all original dataset columns and order.
- Reads the ATK canonical runnable dataset maintained under .atk/datasets/.
- Requires the ATK canonical dataset identity column: atk_id.
- Appends required actual-output column: agent_output.
- Writes results incrementally to .atk/results/vN/eval_results.csv.
- Supports bounded runs with --limit and --offset.
- Supports concurrent runs with --concurrency.
- Supports rerunning only the latest failure_cases.csv samples with --only-failures.
- Optionally writes .atk/results/vN/app.log when log capture is configured.
"""

from __future__ import annotations

import argparse
import contextvars
import csv
import json
import logging
import os
import sys
import threading
import traceback
from collections.abc import Iterator
from concurrent.futures import ThreadPoolExecutor, as_completed
from contextlib import contextmanager, redirect_stderr, redirect_stdout
from pathlib import Path
from typing import Any

RESULTS_DIR = Path(".atk/results")
RUNNER_DIR = Path(".atk/runner")
DATASETS_DIR = Path(".atk/datasets")
DATASET_PATH = DATASETS_DIR / "dataset.csv"
RESULTS_FILENAME = "eval_results.csv"
FAILURE_FILENAME = "failure_cases.csv"
APP_LOG_FILENAME = "app.log"
ATK_ID_FIELD = "atk_id"
PYTHON_LOGGING_CAPTURE_ENABLED = False
CONCURRENT_ROW_LOGGING_ENABLED = True
ROW_LOGGER_NAMES: tuple[str, ...] = ("",)  # Empty string means the root logger.
ROW_LOG_FORMAT = "%(asctime)s %(levelname)s %(name)s: %(message)s"
ROW_LOG_LEVEL = logging.INFO
ROW_LOGS_DIRNAME = "logs"
AGENT_OUTPUT_LOG_PATH_FIELD = "log_path"
ROW_LOGGING_CONCURRENCY_DOWNGRADE_MESSAGE = (
    "Concurrent row-level Python logging capture is disabled; "
    "no per-row log files will be created for --concurrency > 1."
)
REPO_ROOT = Path(__file__).resolve().parents[2]
_ACTIVE_ROW_LOG_CONTEXT: contextvars.ContextVar[dict[str, Any] | None] = contextvars.ContextVar(
    "atk_active_row_log_context",
    default=None,
)

# TODO_AGENT_TUNING_IMPORT_PATHS:
# Add repository import roots required by the target Agent before importing it.
# Common Python src-layout projects need REPO_ROOT / "src"; package-at-root
# projects need REPO_ROOT. Keep this explicit and project-local.
for import_path in (REPO_ROOT / "src", REPO_ROOT):
    import_path_str = str(import_path)
    if import_path.exists() and import_path_str not in sys.path:
        sys.path.insert(0, import_path_str)


class UserActionRequired(RuntimeError):
    """Raised when the user must repair inputs or confirm an unsafe inference."""


class AgentExecutionError(RuntimeError):
    """Raised for a known per-row Agent execution failure that should be recorded."""


def list_version_dirs(results_dir: Path = RESULTS_DIR) -> list[tuple[int, Path]]:
    # Return [(number, path)] for directories named vN where N is a positive integer.
    # Missing results_dir means no versions exist yet; runner first run will create v1.
    if not results_dir.exists():
        return []
    return sorted(
        (int(path.name[1:]), path)
        for path in results_dir.iterdir()
        if path.is_dir() and path.name.startswith("v") and path.name[1:].isdigit() and int(path.name[1:]) > 0
    )


def allocate_next_results_version(results_dir: Path = RESULTS_DIR) -> Path:
    # Used only by eval_runner.py. Current max version determines create vs reuse.
    versions = list_version_dirs(results_dir)
    if not versions:
        target = results_dir / "v1"
    else:
        max_n, current = versions[-1]
        target = results_dir / f"v{max_n + 1}" if (current / "eval_results.csv").exists() else current
    target.mkdir(parents=True, exist_ok=True)
    return target


def resolve_latest_failure_cases_version(results_dir: Path = RESULTS_DIR) -> Path:
    """Return the newest results version that contains failure_cases.csv."""
    for _, version_dir in reversed(list_version_dirs(results_dir)):
        if (version_dir / FAILURE_FILENAME).exists():
            return version_dir
    raise UserActionRequired(
        f"No prior {FAILURE_FILENAME} found; run atk-diagnose first."
    )


def load_dataset(path: Path) -> tuple[list[str], list[dict[str, str]]]:
    """Load CSV rows while preserving original columns.

    TODO_AGENT_TUNING: adapt this function only if the dataset is not CSV.
    Do not rename user columns here. If the source dataset contains agent_output
    or agent_output_log_path, the Skill must confirm a rename strategy before
    generating this script.
    """
    if not path.exists():
        raise UserActionRequired(f"Dataset not found: {path}")
    with path.open("r", newline="", encoding="utf-8-sig") as handle:
        reader = csv.DictReader(handle)
        if not reader.fieldnames:
            raise UserActionRequired(f"Dataset has no header row: {path}")
        fieldnames = list(reader.fieldnames)
        if ATK_ID_FIELD not in fieldnames:
            raise UserActionRequired(
                f"Dataset is missing required {ATK_ID_FIELD} column; rerun atk-init to rebuild "
                "the ATK canonical dataset with 1-based source row IDs."
            )
        rows: list[dict[str, str]] = []
        seen_atk_ids: set[str] = set()
        for row_number, row in enumerate(reader, start=1):
            atk_id = row.get(ATK_ID_FIELD, "").strip()
            if not atk_id:
                raise UserActionRequired(f"Dataset row {row_number} has an empty {ATK_ID_FIELD} value.")
            if not atk_id.isdigit() or int(atk_id) < 1:
                raise UserActionRequired(
                    f"Dataset row {row_number} has invalid {ATK_ID_FIELD}={atk_id!r}; expected a positive integer."
                )
            if atk_id in seen_atk_ids:
                raise UserActionRequired(f"Dataset has duplicate {ATK_ID_FIELD} value: {atk_id}")
            seen_atk_ids.add(atk_id)
            rows.append(row)

        reserved_output_fields = {"agent_output", AGENT_OUTPUT_LOG_PATH_FIELD}
        conflicting_fields = [field for field in fieldnames if field in reserved_output_fields]
        if conflicting_fields:
            raise UserActionRequired(
                "Dataset already contains reserved output column(s) "
                f"{', '.join(conflicting_fields)}; confirm a rename strategy before running."
            )
        return fieldnames, rows


def load_failure_atk_ids(version_dir: Path) -> set[str]:
    """Load failed sample IDs from a previous failure_cases.csv."""
    failure_path = version_dir / FAILURE_FILENAME
    with failure_path.open("r", newline="", encoding="utf-8-sig") as handle:
        reader = csv.DictReader(handle)
        if not reader.fieldnames:
            raise UserActionRequired(f"{failure_path} has no header row.")
        if ATK_ID_FIELD not in reader.fieldnames:
            raise UserActionRequired(f"{failure_path} is missing required {ATK_ID_FIELD} column.")
        failure_ids: set[str] = set()
        for row_number, row in enumerate(reader, start=1):
            atk_id = row.get(ATK_ID_FIELD, "").strip()
            if not atk_id:
                raise UserActionRequired(f"{failure_path} row {row_number} has an empty {ATK_ID_FIELD} value.")
            failure_ids.add(atk_id)
    if not failure_ids:
        raise UserActionRequired(f"{failure_path} contains no failure rows to rerun.")
    return failure_ids


def build_agent_input(row: dict[str, str]) -> Any:
    """Convert a dataset row into the local Agent input.

    TODO_AGENT_TUNING: fill this after inspecting Agent source and dataset columns.
    Example patterns:
    - return row["question"]
    - return {"messages": [{"role": "user", "content": row["query"]}]}
    - return {"input": row["input"], "context": row.get("context", "")}
    """
    raise UserActionRequired("TODO_AGENT_TUNING: implement build_agent_input(row).")


def call_agent(agent_input: Any) -> Any:
    """Call the local Agent and return its raw output.

    TODO_AGENT_TUNING: import and call the target Agent here. Prefer the target
    project's own virtualenv/dependencies. Handle async explicitly if needed.
    Raise AgentExecutionError only for known per-row failures that should be
    preserved in eval_results.csv; let configuration/import errors fail the run.
    """
    raise UserActionRequired("TODO_AGENT_TUNING: implement call_agent(agent_input).")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run Agent tuning cases.")
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Maximum number of dataset rows to run after applying --offset.",
    )
    parser.add_argument(
        "--offset",
        type=int,
        default=0,
        help="Number of dataset rows to skip before running.",
    )
    parser.add_argument(
        "--no-progress",
        action="store_true",
        help="Suppress per-row progress output.",
    )
    parser.add_argument(
        "--concurrency",
        type=int,
        default=1,
        help=(
            "Maximum number of rows to execute concurrently. Use 1 for strict serial "
            "execution; values greater than 1 write rows as they complete."
        ),
    )
    parser.add_argument(
        "--only-failures",
        action="store_true",
        help=(
            "Run only rows whose atk_id appears in the latest prior failure_cases.csv. "
            "--offset and --limit apply after this failure subset is selected."
        ),
    )
    args = parser.parse_args()
    if args.limit is not None and args.limit < 0:
        raise UserActionRequired("--limit must be greater than or equal to 0.")
    if args.offset < 0:
        raise UserActionRequired("--offset must be greater than or equal to 0.")
    if args.concurrency < 1:
        raise UserActionRequired("--concurrency must be greater than or equal to 1.")
    return args


def select_dataset_rows(
    rows: list[dict[str, str]],
    *,
    offset: int,
    limit: int | None,
) -> list[tuple[int, dict[str, str]]]:
    indexed_rows = list(enumerate(rows, start=1))
    selected = indexed_rows[offset:]
    if limit is not None:
        selected = selected[:limit]
    return selected


def select_failure_dataset_rows(
    rows: list[dict[str, str]],
    failure_atk_ids: set[str],
    *,
    offset: int,
    limit: int | None,
) -> list[tuple[int, dict[str, str]]]:
    indexed_rows = list(enumerate(rows, start=1))
    matched = [(source_index, row) for source_index, row in indexed_rows if row[ATK_ID_FIELD] in failure_atk_ids]
    matched_ids = {row[ATK_ID_FIELD] for _, row in matched}
    missing_ids = sorted(
        failure_atk_ids - matched_ids,
        key=lambda value: (0, int(value)) if value.isdigit() else (1, value),
    )
    if missing_ids:
        raise UserActionRequired(
            "Failure atk_id values are missing from the canonical dataset: " + ", ".join(missing_ids)
        )
    selected = matched[offset:]
    if limit is not None:
        selected = selected[:limit]
    return selected


def normalize_agent_output(raw_output: Any) -> str:
    """Serialize Agent output for the required agent_output column."""
    if isinstance(raw_output, str):
        return raw_output
    return json.dumps(raw_output, ensure_ascii=False, sort_keys=True)


def row_log_relative_path(source_index: int) -> Path:
    """Return the stable source-row-numbered path stored in eval_results.csv."""
    return Path(ROW_LOGS_DIRNAME) / f"row_{source_index:06d}.log"


def _configured_row_loggers() -> list[logging.Logger]:
    """Resolve configured Python loggers without mutating project handlers."""
    loggers: list[logging.Logger] = []
    seen: set[int] = set()
    for logger_name in ROW_LOGGER_NAMES:
        logger = logging.getLogger(logger_name) if logger_name else logging.getLogger()
        logger_id = id(logger)
        if logger_id not in seen:
            loggers.append(logger)
            seen.add(logger_id)
    return loggers


class ATKRowLogHandler(logging.Handler):
    """ATK-owned handler that routes records to the active row log only."""

    def __init__(self) -> None:
        super().__init__(ROW_LOG_LEVEL)
        self.setFormatter(logging.Formatter(ROW_LOG_FORMAT))
        self._lock = threading.RLock()
        self._handles: dict[Path, Any] = {}

    def emit(self, record: logging.LogRecord) -> None:
        row_context = _ACTIVE_ROW_LOG_CONTEXT.get()
        if row_context is None:
            return
        log_path = row_context.get("log_path")
        if not isinstance(log_path, Path):
            return

        if getattr(record, "_atk_row_log_written", False):
            return
        setattr(record, "_atk_row_log_written", True)
        try:
            message = self.format(record)
            with self._lock:
                handle = self._handles.get(log_path)
                if handle is None:
                    log_path.parent.mkdir(parents=True, exist_ok=True)
                    handle = log_path.open("a", encoding="utf-8")
                    self._handles[log_path] = handle
                handle.write(message + "\n")
                handle.flush()
        except (OSError, RuntimeError):
            self.handleError(record)

    def close(self) -> None:
        with self._lock:
            for handle in self._handles.values():
                handle.close()
            self._handles.clear()
        super().close()


@contextmanager
def capture_python_row_logging(*, enabled: bool) -> Iterator[None]:
    """Install one run-scoped ATK row-log router for configured loggers."""
    if not enabled:
        yield
        return

    row_handler = ATKRowLogHandler()
    owned_loggers = _configured_row_loggers()
    previous_levels: dict[logging.Logger, int] = {}
    try:
        for logger in owned_loggers:
            previous_levels[logger] = logger.level
            logger.setLevel(ROW_LOG_LEVEL)
            logger.addHandler(row_handler)
        yield
    finally:
        for logger in owned_loggers:
            logger.removeHandler(row_handler)
            logger.setLevel(previous_levels[logger])
        row_handler.close()


@contextmanager
def capture_python_app_logging(version_dir: Path, *, enabled: bool) -> Iterator[None]:
    """Install one run-scoped app.log handler for configured Python loggers."""
    if not enabled:
        yield
        return

    log_path = version_dir / APP_LOG_FILENAME
    log_path.parent.mkdir(parents=True, exist_ok=True)
    app_handler = logging.FileHandler(log_path, mode="w", encoding="utf-8")
    app_handler.setLevel(ROW_LOG_LEVEL)
    app_handler.setFormatter(logging.Formatter(ROW_LOG_FORMAT))
    owned_loggers = _configured_row_loggers()
    previous_levels: dict[logging.Logger, int] = {}
    try:
        for logger in owned_loggers:
            previous_levels[logger] = logger.level
            logger.setLevel(ROW_LOG_LEVEL)
            logger.addHandler(app_handler)
        yield
    finally:
        for logger in owned_loggers:
            logger.removeHandler(app_handler)
            logger.setLevel(previous_levels[logger])
        app_handler.close()


@contextmanager
def capture_row_python_logs(
    version_dir: Path,
    source_index: int,
    *,
    enabled: bool,
) -> Iterator[str]:
    """Set ATK row context for one row and create its referenced log file."""
    if not enabled:
        yield ""
        return

    relative_path = row_log_relative_path(source_index)
    log_path = version_dir / relative_path
    log_path.parent.mkdir(parents=True, exist_ok=True)
    log_path.open("w", encoding="utf-8").close()

    token = _ACTIVE_ROW_LOG_CONTEXT.set({"relative_path": relative_path, "log_path": log_path})
    try:
        yield relative_path.as_posix()
    finally:
        _ACTIVE_ROW_LOG_CONTEXT.reset(token)


def build_result_row(
    index: int,
    row: dict[str, str],
    *,
    version_dir: Path,
    row_logging_enabled: bool,
) -> dict[str, str]:
    output_row = dict(row)
    output_row[AGENT_OUTPUT_LOG_PATH_FIELD] = ""
    with capture_row_python_logs(version_dir, index, enabled=row_logging_enabled) as row_log_path:
        if row_log_path:
            output_row[AGENT_OUTPUT_LOG_PATH_FIELD] = row_log_path
        try:
            agent_input = build_agent_input(row)
            raw_output = call_agent(agent_input)
            output_row["agent_output"] = normalize_agent_output(raw_output)
            output_row["output_status"] = "ok"
        except UserActionRequired:
            # Configuration/TODO/confirmation failures must stop the run; otherwise
            # an invalid runner could create a misleading successful eval_results.csv.
            raise
        except AgentExecutionError as exc:
            # Known per-row Agent failures are still useful evaluation evidence.
            output_row["agent_output"] = json.dumps(
                {
                    "error": type(exc).__name__,
                    "message": str(exc),
                    "row_number": index,
                    "traceback": traceback.format_exc(limit=8),
                },
                ensure_ascii=False,
            )
            output_row["output_status"] = "error"
            output_row["error_type"] = type(exc).__name__
            output_row["error_message"] = str(exc)
    return output_row


def get_output_fieldnames(
    source_fieldnames: list[str],
    rows: list[dict[str, str]],
) -> list[str]:
    fixed_output_fields = [
        "output_status",
        "error_type",
        "error_message",
        AGENT_OUTPUT_LOG_PATH_FIELD,
    ]
    auxiliary_fields = sorted(
        {
            key
            for row in rows
            for key in row
            if key.startswith("agent_output_")
            and key not in source_fieldnames
            and key not in fixed_output_fields
        }
    )
    return list(source_fieldnames) + ["agent_output"] + fixed_output_fields + auxiliary_fields


def open_results_writer(
    version_dir: Path,
    source_fieldnames: list[str],
    rows: list[dict[str, str]],
) -> tuple[Path, Any, csv.DictWriter[str]]:
    output_path = version_dir / RESULTS_FILENAME
    fieldnames = get_output_fieldnames(source_fieldnames, rows)
    handle = output_path.open("w", newline="", encoding="utf-8")
    writer = csv.DictWriter(handle, fieldnames=fieldnames, extrasaction="ignore")
    writer.writeheader()
    handle.flush()
    os.fsync(handle.fileno())
    return output_path, handle, writer


def run_rows(
    version_dir: Path,
    source_fieldnames: list[str],
    selected_rows: list[tuple[int, dict[str, str]]],
    *,
    total_rows: int,
    show_progress: bool,
    concurrency: int,
    row_logging_enabled: bool,
) -> Path:
    output_path, handle, writer = open_results_writer(
        version_dir,
        source_fieldnames,
        [row for _, row in selected_rows],
    )
    try:
        with capture_python_row_logging(enabled=row_logging_enabled):
            for completed, (source_index, result_row) in enumerate(
                execute_rows(
                    selected_rows,
                    version_dir=version_dir,
                    concurrency=concurrency,
                    row_logging_enabled=row_logging_enabled,
                ),
                start=1,
            ):
                writer.writerow(result_row)
                handle.flush()
                os.fsync(handle.fileno())
                if show_progress:
                    print(
                        f"Processed row {source_index}/{total_rows} "
                        f"({completed}/{len(selected_rows)} selected)",
                        flush=True,
                    )
    finally:
        handle.close()
    return output_path


def execute_rows(
    selected_rows: list[tuple[int, dict[str, str]]],
    *,
    version_dir: Path,
    concurrency: int,
    row_logging_enabled: bool,
) -> Iterator[tuple[int, dict[str, str]]]:
    """Execute rows and return results as they complete.

    With concurrency=1 this preserves dataset order. With concurrency>1, rows are
    yielded in completion order so eval_results.csv remains useful during long runs.
    """
    if concurrency == 1:
        for source_index, row in selected_rows:
            yield source_index, build_result_row(
                source_index,
                row,
                version_dir=version_dir,
                row_logging_enabled=row_logging_enabled,
            )
        return

    with ThreadPoolExecutor(max_workers=concurrency) as executor:
        futures = {
            executor.submit(
                build_result_row,
                source_index,
                row,
                version_dir=version_dir,
                row_logging_enabled=row_logging_enabled,
            ): source_index
            for source_index, row in selected_rows
        }
        for future in as_completed(futures):
            source_index = futures[future]
            yield source_index, future.result()


def run() -> int:
    args = parse_args()
    failure_atk_ids: set[str] | None = None
    if args.only_failures:
        failure_version_dir = resolve_latest_failure_cases_version()
        failure_atk_ids = load_failure_atk_ids(failure_version_dir)
    version_dir = allocate_next_results_version()
    source_fieldnames, dataset_rows = load_dataset(DATASET_PATH)
    if args.only_failures:
        if failure_atk_ids is None:
            raise UserActionRequired("No failure atk_id values were loaded.")
        selected_rows = select_failure_dataset_rows(
            dataset_rows,
            failure_atk_ids,
            offset=args.offset,
            limit=args.limit,
        )
    else:
        selected_rows = select_dataset_rows(dataset_rows, offset=args.offset, limit=args.limit)
    row_logging_enabled = PYTHON_LOGGING_CAPTURE_ENABLED and (
        args.concurrency == 1 or CONCURRENT_ROW_LOGGING_ENABLED
    )
    if PYTHON_LOGGING_CAPTURE_ENABLED and args.concurrency > 1 and not CONCURRENT_ROW_LOGGING_ENABLED:
        print(ROW_LOGGING_CONCURRENCY_DOWNGRADE_MESSAGE, file=sys.__stderr__, flush=True)
    with capture_python_app_logging(version_dir, enabled=PYTHON_LOGGING_CAPTURE_ENABLED):
        output_path = run_rows(
            version_dir,
            source_fieldnames,
            selected_rows,
            total_rows=len(dataset_rows),
            show_progress=not args.no_progress,
            concurrency=args.concurrency,
            row_logging_enabled=row_logging_enabled,
        )
    print(f"Wrote {len(selected_rows)} rows to {output_path}")
    return 0


if __name__ == "__main__":
    # TODO_AGENT_TUNING_LOG_CAPTURE:
    # If configured Python logging is available, run() writes app.log through logging.
    # If the Agent emits useful stdout/stderr logs instead, keep this redirect fallback.
    # If the Agent writes its own log file, replace this with explicit copy/read logic.
    # If no reliable log source exists, call sys.exit(run()) directly and omit app.log.
    target_version_dir = allocate_next_results_version()

    # Reuse the preallocated version directory for this process so log and results align.
    def allocate_next_results_version(results_dir=RESULTS_DIR):  # type: ignore[no-redef]
        target_version_dir.mkdir(parents=True, exist_ok=True)
        return target_version_dir

    if PYTHON_LOGGING_CAPTURE_ENABLED:
        sys.exit(run())

    log_path = target_version_dir / APP_LOG_FILENAME
    with log_path.open("w", encoding="utf-8") as log_handle:
        with redirect_stdout(log_handle), redirect_stderr(log_handle):
            exit_code = run()
    print(f"Wrote optional stdout/stderr app log to {log_path}")
    sys.exit(exit_code)