# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
Data models for the Episteme Environment.

Episteme trains LLMs to investigate before they conclude — rewarding careful,
grounded reasoning over confident hallucination. These models define the
action, observation, and state schemas for the environment.
"""

from typing import Any, Dict, List, Literal

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


class EpistemeAction(Action):
    tool_name: Literal[  #tools given to the agent
        "search_corpus",
        "read_document",
        "write_note",
        "read_notes",
        "cross_check",
        "flag_uncertainty",
        "file_output",
    ] = Field(
        ...,
        description="Name of the research tool to invoke.",
    )
    arguments: Dict[str, Any] = Field(
        default_factory=dict,
        description="Tool-specific keyword arguments.",
    )


class EpistemeObservation(Observation):

    result: str = Field(
        default="",
        description="Textual result produced by the tool that was invoked.",
    )
    reward: float = Field(
        default=0.0,
        description="Reward signal for this step.",
    )
    done: bool = Field(
        default=False,
        description="Whether the episode has terminated.",
    )
    step_count: int = Field(
        default=0,
        ge=0,
        description="Number of steps taken so far in this episode.",
    )
    steps_remaining: int = Field(
        default=0,
        ge=0,
        description="Number of steps the agent has left before the budget expires.",
    )
    success: bool = Field(
        default=True,
        description="Whether the tool invocation succeeded.",
    )


class EpistemeState(State):
    episode_id: str = Field(
        default="",
        description="Unique identifier for the current episode.",
    )
    task_type: str = Field(
        ...,
        description="Category of the research task (e.g. 'fact_check', 'synthesis').",
    )
    task_description: str = Field(
        ...,
        description="Natural-language description of what the agent must investigate.",
    )
    step_budget: int = Field(
        default=20,
        ge=1,
        description="Maximum number of steps allowed in this episode.",
    )
    step_count: int = Field(
        default=0,
        ge=0,
        description="Number of steps taken so far.",
    )
    notes: Dict[str, Any] = Field(
        default_factory=dict,
        description="Agent's scratchpad — free-form key-value store for notes.",
    )
    corpus: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="List of document dicts available for this episode.",
    )
    done: bool = Field(
        default=False,
        description="Whether the episode has ended.",
    )
    final_reward: float = Field(
        default=0.0,
        description="Cumulative or final reward for the episode.",
    )
    noise_reads: int = Field(
        default=0,
        ge=0,
        description="Number of noise documents read this episode.",
    )