"""Typed OpenEnv wire models for SpreadsheetBench."""

from __future__ import annotations

from typing import Any

from openenv.core.env_server.types import Action, Observation, State
from pydantic import BaseModel, ConfigDict, Field, model_validator

from spreadsheetbench_env.interface import STEP_LIMIT


class SpreadsheetBenchSubmitScript(BaseModel):
    """Submit one complete Python script for SpreadsheetBench."""

    model_config = ConfigDict(extra="forbid")

    code: str = Field(
        min_length=1,
        description=(
            "One complete Python script. Read the injected INPUT_PATH variable and "
            "write the modified workbook to the injected OUTPUT_PATH variable. A failing "
            "execution is returned as feedback for a corrected script; the final script is "
            "run independently on all hidden cases."
        ),
    )


class SpreadsheetBenchAction(Action, SpreadsheetBenchSubmitScript):
    """One complete Python script submitted on one repair turn."""

    @model_validator(mode="before")
    @classmethod
    def _accept_native_envelope(cls, values: Any) -> Any:
        if not isinstance(values, dict):
            return values
        native_keys = {"name", "arguments"}
        if not native_keys.intersection(values):
            return values
        if set(values) != native_keys:
            raise ValueError(
                "SpreadsheetBench native tool envelope must contain exactly "
                "'name' and 'arguments'"
            )
        action = native_tool_action(values["name"], values["arguments"])
        return action.model_dump()


def native_tool_specs() -> tuple[dict[str, Any], ...]:
    """Return the deterministic model-facing native tool catalog."""

    return (
        {
            "name": "submit_script",
            "description": SpreadsheetBenchSubmitScript.__doc__,
            "input_schema": SpreadsheetBenchSubmitScript.model_json_schema(),
        },
    )


def native_tool_action(name: str, arguments: dict[str, Any]) -> SpreadsheetBenchAction:
    """Validate one native tool call and convert it to the OpenEnv action."""

    if name != "submit_script":
        raise ValueError(f"unknown SpreadsheetBench native tool: {name!r}")
    native_action = SpreadsheetBenchSubmitScript.model_validate(arguments)
    return SpreadsheetBenchAction(**native_action.model_dump())


class SpreadsheetBenchObservation(Observation):
    """Actor-visible task or terminal summary without hidden grading evidence."""

    task_id: int
    text: str
    goal: str
    task_type: str
    case_count: int = Field(ge=1)
    answer_position: str
    success: bool
    progress: float = Field(ge=0.0, le=1.0)
    terminated: bool
    truncated: bool
    step_limit: int = STEP_LIMIT

    @model_validator(mode="before")
    @classmethod
    def _derive_done(cls, values: Any) -> Any:
        if isinstance(values, dict):
            values = dict(values)
            values["done"] = bool(
                values.get("terminated", False) or values.get("truncated", False)
            )
        return values


class SpreadsheetBenchState(State):
    """Evaluator-facing state for one connection-owned episode."""

    task_id: int | None = None
    goal: str | None = None
    reward: float = 0.0
    success: bool = False
    progress: float = Field(default=0.0, ge=0.0, le=1.0)
    terminated: bool = False
    truncated: bool = False
    step_limit: int = STEP_LIMIT
    diagnostics: dict[str, str] = Field(default_factory=dict)
    closed: bool = False