"""Framework-neutral SpreadsheetBench episode lifecycle: SkillOpt's multi-turn code repair.

Mirrors ``skillopt.envs.spreadsheetbench.codegen_agent`` in ``mode=multi``: every turn the agent
submits one complete script, the harness executes it on the first hidden case, and a failing
execution comes back as feedback for another turn, up to 30 turns within a 600 s task budget.
The episode ends on the first successful execution (evaluation) or, during training, on the first
execution whose first-case output grades correct; a wrong output is fed back as SkillOpt's
cell-level report with expected values hidden. The final script is then graded independently on
every hidden case, as SkillOpt's ``process_one`` does.
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
import tempfile
import time
from typing import Protocol
from uuid import uuid4

from spreadsheetbench_env.catalog import SpreadsheetTask, get_task
from spreadsheetbench_env.executor import ExecutionResult, ScriptExecutor
from spreadsheetbench_env.grader import InfraError, compare_workbooks, verification_report
from spreadsheetbench_env.interface import (
    DEADLINE_MARGIN_S,
    STEP_LIMIT,
    TASK_TIME_BUDGET_S,
    build_task_prompt,
    eval_feedback,
    execution_error_feedback,
    extract_script,
    preview_workbook,
)


@dataclass(frozen=True, slots=True)
class CoreObservation:
    task_id: int
    text: str
    goal: str
    task_type: str
    case_count: int
    answer_position: str
    success: bool
    progress: float
    step_limit: int = STEP_LIMIT


@dataclass(frozen=True, slots=True)
class CoreStep:
    observation: CoreObservation
    reward: float
    terminated: bool
    truncated: bool


@dataclass(frozen=True, slots=True)
class CoreState:
    task_id: int | None
    goal: str | None
    reward: float
    success: bool
    progress: float
    terminated: bool
    truncated: bool
    step_count: int
    step_limit: int
    diagnostics: dict[str, str]
    closed: bool


class _SessionExecutor(Protocol):
    @property
    def root(self) -> Path: ...

    def run(
        self,
        code: str,
        input_path: str | Path,
        *,
        case_name: str,
    ) -> ExecutionResult: ...

    def remove_case(self, case_name: str) -> None: ...

    def close(self) -> None: ...


_TaskLookup = Callable[[int], SpreadsheetTask]
_ExecutorFactory = Callable[[Path], _SessionExecutor]


def _default_executor_factory(root: Path) -> _SessionExecutor:
    return ScriptExecutor(root)


class SpreadsheetBenchCore:
    """Own one isolated multi-turn SpreadsheetBench episode."""

    def __init__(
        self,
        data_root: str | Path,
        *,
        task_lookup: _TaskLookup = get_task,
        executor_factory: _ExecutorFactory = _default_executor_factory,
        work_root: str | Path | None = None,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        self._data_root = Path(data_root).resolve()
        if not self._data_root.is_dir():
            raise FileNotFoundError(
                f"SpreadsheetBench runtime data root does not exist: {self._data_root}"
            )
        self._task_lookup = task_lookup
        self._executor_factory = executor_factory
        self._work_root = Path(work_root or tempfile.gettempdir()).resolve()
        self._clock = clock
        self._executor: _SessionExecutor | None = None
        self._task: SpreadsheetTask | None = None
        self._cases: tuple[tuple[Path, Path], ...] = ()
        self._prompt = ""
        self._preview = ""
        self._training_feedback = False
        self._deadline = 0.0
        self._last_code = ""
        self._reward = 0.0
        self._progress = 0.0
        self._terminated = False
        self._truncated = False
        self._step_count = 0
        self._diagnostics: dict[str, str] = {}
        self._closed = False

    def _runtime_path(self, relative: str, *, hidden_answer: bool) -> Path:
        candidate = (self._data_root / relative).resolve()
        if not candidate.is_relative_to(self._data_root):
            raise InfraError(f"runtime path escapes data root: {relative!r}")
        if not candidate.is_file():
            kind = "hidden answer" if hidden_answer else "input workbook"
            raise InfraError(f"{kind} is missing: {relative}")
        return candidate

    def reset(self, task_id: int, *, training_feedback: bool) -> CoreObservation:
        """Start ``task_id``. ``training_feedback`` enables SkillOpt's training-only gold check:
        a successfully executed but wrong first-case output is fed back instead of ending the
        episode. Evaluation must pass ``False`` so no grading evidence reaches the agent."""
        self._ensure_open()
        if not isinstance(training_feedback, bool):
            raise TypeError("training_feedback must be a bool")
        if self._executor is not None:
            self._executor.close()
        task = self._task_lookup(int(task_id))
        cases = tuple(
            (
                self._runtime_path(case.input_path, hidden_answer=False),
                self._runtime_path(case.answer_path, hidden_answer=True),
            )
            for case in task.cases
        )
        if not cases:
            raise InfraError(f"task {task.task_id} has no runtime cases")
        session_root = self._work_root / f"spreadsheetbench-session-{uuid4().hex}"
        executor = self._executor_factory(session_root)
        try:
            preview = preview_workbook(cases[0][0])
        except Exception as error:
            executor.close()
            raise InfraError(
                f"could not preview input workbook for task {task.task_id}: {error}"
            ) from error

        self._executor = executor
        self._task = task
        self._cases = cases
        self._preview = preview
        self._prompt = build_task_prompt(
            task.instruction,
            task.task_type,
            task.answer_position,
            preview,
        )
        self._training_feedback = training_feedback
        self._deadline = self._clock() + TASK_TIME_BUDGET_S
        self._last_code = ""
        self._reward = 0.0
        self._progress = 0.0
        self._terminated = False
        self._truncated = False
        self._step_count = 0
        self._diagnostics = {}
        return self._observation(self._prompt)

    def step(self, action_code: str) -> CoreStep:
        self._ensure_open()
        if self._task is None or self._executor is None or not self._cases:
            raise RuntimeError("SpreadsheetBench core must be reset before step")
        if self._terminated or self._truncated:
            raise RuntimeError("SpreadsheetBench episode has ended")

        turn = self._step_count + 1
        self._step_count = turn
        if self._deadline - self._clock() <= DEADLINE_MARGIN_S:
            # SkillOpt stops issuing turns here, so the script produced on this turn never
            # existed for it: the previous script is the one that gets graded.
            return self._finish(
                self._last_code,
                truncated=True,
                preface=(
                    f"The {TASK_TIME_BUDGET_S:g} s task time budget was exhausted before this "
                    "turn; the previous script is graded."
                ),
            )

        code = extract_script(action_code)
        self._last_code = code
        first_input, first_answer = self._cases[0]
        case_name = f"turn-{turn}"
        try:
            execution = self._executor.run(code, first_input, case_name=case_name)
            if execution.ok and execution.output_path is not None:
                if not self._training_feedback:
                    return self._finish(code, truncated=False)
                matched, _ = compare_workbooks(
                    first_answer, execution.output_path, self._task.answer_position
                )
                if matched:
                    return self._finish(code, truncated=False)
                feedback = eval_feedback(
                    verification_report(
                        execution.output_path, first_answer, self._task.answer_position
                    )
                )
            else:
                feedback = execution_error_feedback(execution.error)
        finally:
            self._executor.remove_case(case_name)

        if turn >= STEP_LIMIT:
            return self._finish(code, truncated=True, preface=feedback)
        return CoreStep(
            observation=self._observation(feedback),
            reward=0.0,
            terminated=False,
            truncated=False,
        )

    def _finish(self, code: str, *, truncated: bool, preface: str = "") -> CoreStep:
        """Grade ``code`` independently on every hidden case and end the episode."""
        assert self._task is not None and self._executor is not None
        summaries: list[str] = []
        verification: list[str] = []
        passed = 0
        fail_reason = ""
        for number, (input_path, answer_path) in enumerate(self._cases, start=1):
            case_name = f"case-{number}"
            try:
                execution = self._executor.run(code, input_path, case_name=case_name)
                if not execution.ok or execution.output_path is None:
                    summaries.append(f"Case {number}: execution ERROR")
                    verification.append(
                        f"## Execution (case {number})\nERROR: {execution.error}"
                    )
                    if not fail_reason:
                        final_line = (execution.error.strip().splitlines() or ["unknown"])[-1]
                        fail_reason = f"exec-error: {final_line[:200]}"
                    if execution.timed_out:
                        break
                    continue

                matched, reason = compare_workbooks(
                    answer_path,
                    execution.output_path,
                    self._task.answer_position,
                )
                summaries.append(
                    f"Case {number}: execution ok; grade {'PASS' if matched else 'FAIL'}"
                )
                if matched:
                    passed += 1
                elif not fail_reason:
                    fail_reason = f"eval-mismatch: {reason[:200]}"
                verification.append(
                    f"## Eval Result (case {number}): {'PASS' if matched else 'FAIL'}\n"
                    f"{reason}\n\n"
                    f"{verification_report(execution.output_path, answer_path, self._task.answer_position)}"
                )
            finally:
                self._executor.remove_case(case_name)

        self._terminated = not truncated
        self._truncated = truncated
        self._progress = passed / len(self._cases)
        self._reward = 1.0 if passed == len(self._cases) else 0.0
        self._diagnostics = {
            "verification": "[POST-EXECUTION VERIFICATION]\n\n"
            + "\n\n---\n\n".join(verification),
            "spreadsheet_preview": self._preview,
            "fail_reason": "" if self._reward == 1.0 else fail_reason,
            "code": code,
            "turns": str(self._step_count),
            "training_feedback": "true" if self._training_feedback else "false",
        }
        text = "\n".join(summaries)
        if preface:
            text = f"{preface}\n\n{text}"
        return CoreStep(
            observation=self._observation(text),
            reward=self._reward,
            terminated=self._terminated,
            truncated=self._truncated,
        )

    def state(self) -> CoreState:
        return CoreState(
            task_id=None if self._task is None else self._task.task_id,
            goal=None if self._task is None else self._task.instruction,
            reward=self._reward,
            success=self._reward == 1.0 and (self._terminated or self._truncated),
            progress=self._progress,
            terminated=self._terminated,
            truncated=self._truncated,
            step_count=self._step_count,
            step_limit=STEP_LIMIT,
            diagnostics=dict(self._diagnostics),
            closed=self._closed,
        )

    def close(self) -> None:
        if self._closed:
            return
        if self._executor is not None:
            self._executor.close()
        self._closed = True

    def _observation(self, text: str) -> CoreObservation:
        if self._task is None:
            raise RuntimeError("SpreadsheetBench core has no current task")
        return CoreObservation(
            task_id=self._task.task_id,
            text=text,
            goal=self._task.instruction,
            task_type=self._task.task_type,
            case_count=len(self._cases),
            answer_position=self._task.answer_position,
            success=self._reward == 1.0 and (self._terminated or self._truncated),
            progress=self._progress,
        )

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("SpreadsheetBench core is closed")