# NOTE : task_id IS THE PIPELINE_ID INTENRALLY

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

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

from cloudenv.server.providers.aws.services import AWSService

################################################################################
# OpenEnv Models
################################################################################


class CloudEnvAction(Action):
    task_id: str = Field(
        ..., min_length=1, description="Task identifier (pipeline_id)"
    )
    episode_id: str = Field(..., min_length=1, description="Episode identifier")

    service: Literal["internal"] | AWSService = Field(
        ..., description="AWS service (s3, ec2, etc)"
    )
    instance_id: Optional[str] = Field(
        None, description="Instance of AWS service (s3, ec2, etc)"
    )
    operation: str = Field(
        ..., description="Operation name (CreateBucket, RunInstances, etc)"
    )
    payload: Dict[str, Any] = Field(
        default_factory=dict, description="Arguments for the AWS API call"
    )
    principal_arn: Optional[str] = Field(
        default=None, description="IAM principal performing the action"
    )

    intent: Optional[str] = Field(
        default=None,
        description="Optional semantic intent (e.g., 'fix_s3_access')",
    )


class CloudEnvState(State):
    task_id: Optional[str] = None
    episode_id: Optional[str] = None
    step_count: int = 0

    reward_history: List[List[float]] = Field(default_factory=list)


class CloudEnvObservation(Observation):
    task_id: Optional[str] = None
    episode_id: Optional[str] = None
    step_count: int = 0

    operation_result: Dict[str, Any] = Field(default_factory=dict)
    reward: float = 0.0
    done: bool = False

    current_pipeline_state: Dict[str, Any] = Field(default_factory=dict)
    model_config = {"arbitrary_types_allowed": True}