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

from enum import Enum
from typing import Any, Dict, Optional

from openenv.core.env_server.types import Action, Observation

from typing import Optional, Literal



class StockExchangeMarket(Enum):
    BSE = "bse"
    NSE = "nse"



class FinenvAction(Action):
    """
    Action model for Finenv environment.

    Supports:
    - init → initialize environment
    - buy → buy shares
    - sell → sell shares
    - hold → no action
    """

    type: str  # "init", "buy", "sell", "hold"

    # Trading fields
    quantity: Optional[int] = 0

    task_type: Optional[Literal["easy", "medium", "hard"]] = None

    # Initialization fields
    stock: Optional[str] = None
    market: Optional[StockExchangeMarket] = None
    initial_cash: Optional[float] = None
    max_steps: Optional[int] = None


class FinenvObservation(Observation):
    """Observation from the Finenv environment - the echoed message."""
    stock: Optional[str] = None
    market: Optional[StockExchangeMarket] = None
    shares: int
    cash: float
    price: float
    reward: str
    done: bool
    metadata: Optional[Dict[str, Any]] = None