env/models.py

from __future__ import annotations

from typing import Literal

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


class LeanMigrateReward(BaseModel):
    score: float
    cumulative_score: float
    tests_passed: int | None
    tests_total: int | None
    proof_compiled: bool | None
    breakdown: dict[str, float]
    feedback: str
    lean_error: str | None


class LeanMigrateAction(Action):
    type: Literal["inspect", "analyze_deps", "run_tests", "submit"]
    function_name: str
    candidate_code: str | None = Field(
        default=None,
        description="Working implementation used by run_tests.",
    )
    target_code: str | None = Field(
        default=None,
        description="Final target-language implementation submitted for verification.",
    )
    lean_proof: str | None = Field(
        default=None,
        description="Proof text required only for proof tasks such as payment_saga.",
    )

    @model_validator(mode="after")
    def _validate_payload(self) -> "LeanMigrateAction":
        if self.type == "run_tests" and not self.candidate_code:
            raise ValueError("run_tests actions require candidate_code")
        if self.type == "submit" and not (self.target_code or self.lean_proof):
            raise ValueError("submit actions require target_code or lean_proof")
        return self


class InspectAction(LeanMigrateAction):
    type: Literal["inspect"]


class AnalyzeDepsAction(LeanMigrateAction):
    type: Literal["analyze_deps"]


class RunTestsAction(LeanMigrateAction):
    type: Literal["run_tests"]
    candidate_code: str


class SubmitAction(LeanMigrateAction):
    type: Literal["submit"]
    target_code: str | None = None
    lean_proof: str | None = None


class LeanMigrateObservation(Observation):
    episode_id: str
    task_id: str
    episode_step: int
    max_steps: int
    source_language: str
    target_language: str
    source_files: list[str]
    verified: list[str]
    failing: dict[str, str]
    remaining: list[str]
    progress: float
    last_action_type: str | None
    last_action_feedback: str | None
    last_step_reward: float | None
    reward_details: LeanMigrateReward | None = None


__all__ = [
    "AnalyzeDepsAction",
    "InspectAction",
    "LeanMigrateAction",
    "LeanMigrateObservation",
    "LeanMigrateReward",
    "RunTestsAction",
    "SubmitAction",
]