# 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 NeuroCaster OpenEnv environment."""

from typing import Any, Dict, List, Literal, Optional

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


NeurocasterStage = Literal[
    "investigator",
    "drafter",
    "verifier_bio",
    "investigate",
    "draft",
    "critic",
    "render",
    "tribe",
    "done",
]

CurriculumLevel = Literal["easy", "medium", "hard"]


class NeurocasterAction(Action):
    """Action submitted by the worker policy for the current DAG stage."""

    stage: NeurocasterStage = Field(..., description="Stage this action is attempting to complete")
    repo_path: Optional[str] = Field(
        default=None,
        description="Repository path under shared_data/repos for investigator indexing",
    )
    query: Optional[str] = Field(default=None, description="Codebase question to investigate")
    context_chunks: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="Investigator context returned from Neo4j/code indexing",
    )
    slide_markdown: Optional[str] = Field(
        default=None,
        description="Marp/Slidev markdown produced by the drafter",
    )
    mermaid_diagrams: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="Mermaid diagram payloads with name/code/embed path fields",
    )
    artifact_paths: Dict[str, str] = Field(
        default_factory=dict,
        description="Paths to generated artifacts inside /app/shared_data",
    )
    code_queries: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="Optional investigator-authored code-index or Cypher queries",
    )
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Optional agent metadata")


class NeurocasterObservation(Observation):
    """Observation returned by NeuroCaster after each deterministic DAG transition."""

    stage: NeurocasterStage = Field(default="investigator", description="Next expected stage")
    message: str = Field(default="", description="Human-readable environment status")
    curriculum_level: CurriculumLevel = Field(default="easy", description="Selected task difficulty")
    task: Dict[str, Any] = Field(default_factory=dict, description="Visible task prompt and repo target")
    artifact_paths: Dict[str, str] = Field(default_factory=dict, description="Episode artifact paths")
    repo_metadata: Dict[str, Any] = Field(default_factory=dict, description="Selected GitHub/local repo metadata")
    complexity_scores: Dict[str, Any] = Field(default_factory=dict, description="Static repo complexity analysis")
    ingestion_status: Dict[str, Any] = Field(default_factory=dict, description="code_indexing_mcp ingestion result")
    query_plan: Dict[str, Any] = Field(default_factory=dict, description="Planned investigation queries")
    query_logs: List[Dict[str, Any]] = Field(default_factory=list, description="Executed code-index/Cypher query audit log")
    rubric: Dict[str, Any] = Field(default_factory=dict, description="Hidden rubric summary and artifact path")
    curriculum_update: Dict[str, Any] = Field(default_factory=dict, description="Curriculum update from episode outcome")
    verifier_results: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="Programmatic and critic verifier results",
    )
    reward_components: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="Dense and terminal reward components applied this step",
    )
    metrics: Dict[str, float] = Field(
        default_factory=dict,
        description="Logged pass rates and biological z-score metrics",
    )
    fatal_error: Optional[str] = Field(default=None, description="Fatal verifier failure, if any")


class NeurocasterState(State):
    """Serializable OpenEnv state for the deterministic NeuroCaster episode."""

    curriculum_level: CurriculumLevel = Field(default="easy")
    current_stage: NeurocasterStage = Field(default="investigator")
    task: Dict[str, Any] = Field(default_factory=dict)
    hidden_criteria: List[str] = Field(default_factory=list)
    artifact_paths: Dict[str, str] = Field(default_factory=dict)
    repo_metadata: Dict[str, Any] = Field(default_factory=dict)
    complexity_scores: Dict[str, Any] = Field(default_factory=dict)
    ingestion_status: Dict[str, Any] = Field(default_factory=dict)
    query_plan: Dict[str, Any] = Field(default_factory=dict)
    query_logs: List[Dict[str, Any]] = Field(default_factory=list)
    rubric: Dict[str, Any] = Field(default_factory=dict)
    curriculum_update: Dict[str, Any] = Field(default_factory=dict)
    stage_history: List[Dict[str, Any]] = Field(default_factory=list)
    verifier_results: List[Dict[str, Any]] = Field(default_factory=list)
    reward_components: List[Dict[str, Any]] = Field(default_factory=list)
    metrics: Dict[str, float] = Field(default_factory=dict)
    cumulative_reward: float = Field(default=0.0)
    done: bool = Field(default=False)