# 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.
"""
Finenv Environment Implementation (Dynamic Stock via First Action)

This environment simulates stock trading.
The stock and market are NOT fixed — they are initialized dynamically
using the first action of type "init".

Actions supported:
- init → initialize stock, market, capital
- buy → buy shares
- sell → sell shares
- hold → do nothing
"""

from uuid import uuid4

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State
try:
    from ..models import FinenvAction, FinenvObservation, StockExchangeMarket
    from .reward import reward_message
    from .stock_price import get_stock_price
except ImportError:
    from models import FinenvAction, FinenvObservation, StockExchangeMarket
    from server.reward import reward_message
    from server.stock_price import get_stock_price


class FinenvEnvironment(Environment):

    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self):
        """
        Initialize environment with default values.
        Actual stock configuration will be done using 'init' action.
        """

        self._state = State(episode_id=str(uuid4()), step_count=0)

        # Default placeholders (will be overwritten)
        self.stock = None
        self.market = None
        self.initial_cash = 10000.0
        self.max_steps = 100
        self.task = "easy"

        self.cash = self.initial_cash
        self.shares = 0

        self.price_history = []
        self.stock_price = 0.0

        # Flag to check if environment is initialized
        self.initialized = False

    def reset(self) -> FinenvObservation:
        """
        Reset environment WITHOUT changing stock.

        Stock will be set using first action (type="init")
        """

        self._state = State(episode_id=str(uuid4()), step_count=0)

        self.cash = self.initial_cash
        self.shares = 0

        return FinenvObservation(
            stock=self.stock,
            market=self.market.value if self.market else None,
            price=self.stock_price,
            shares=self.shares,
            cash=self.cash,
            done=False,
            reward=reward_message(0),
            metadata={"message": "Send init action to start trading"}
        )

    def step(self, action: FinenvAction) -> FinenvObservation:
        """
        Execute one step.

        First action MUST be:
        {
            "type": "init",
            "stock": "RELIANCE",
            "market": "NSE"
        }
        """

        
        if action.type == "init":
            if not action.stock or not action.market:
                raise ValueError("Stock and market required for init")
            
            self.stock = action.stock
            market=action.market
            if isinstance(market, str):
                market = market.lower()
                if market == "nse":
                    self.market = StockExchangeMarket.NSE
                elif market == "bse":
                    self.market = StockExchangeMarket.BSE
                else:
                    raise ValueError("Unsupported market. Use 'NSE' or 'BSE'.")
            elif isinstance(market, StockExchangeMarket):
                self.market = market
            else:
                raise ValueError("Invalid market type. Use string or StockExchangeMarket enum.")

            

            self.initial_cash = action.initial_cash or 10000.0
            self.max_steps = action.max_steps or 100
            self.task = action.task_type.lower() if action.task_type else "easy"

            self.cash = self.initial_cash
            self.shares = 0

            # Fetch stock data
            self.price_history = get_stock_price(
                self.stock,
                self.market.value,
                self.max_steps
            )

            if not self.price_history:
                raise ValueError("Invalid stock or no data available")

            self.stock_price = self.price_history[0]

            self.initialized = True

            return FinenvObservation(
                stock=self.stock,
                market=self.market.value,
                price=self.stock_price,
                shares=self.shares,
                cash=self.cash,
                done=False,
                reward=reward_message(0),
                metadata={"message": "Environment initialized"}
            )

        
        if not self.initialized:
            raise ValueError("Environment not initialized. Send init action first.")

        # Increment step count
        self._state.step_count += 1

        
        old_value = self.cash + self.shares * self.stock_price

        
        if action.type == "buy":
            cost = action.quantity * self.stock_price
            if self.cash >= cost:
                self.cash -= cost
                self.shares += action.quantity

        elif action.type == "sell":
            if self.shares >= action.quantity:
                self.cash += action.quantity * self.stock_price
                self.shares -= action.quantity

        # hold → no action

       
        index = min(self._state.step_count, len(self.price_history) - 1)
        self.stock_price = self.price_history[index]

        new_value = self.cash + self.shares * self.stock_price
        reward = (new_value - old_value) / self.initial_cash
        reward = max(min(reward, 1.0), -1.0)
        reward=float(reward)
       
        done = self._state.step_count >= self.max_steps
        if done:
            final_value = self.cash + self.shares * self.stock_price
            profit = final_value - self.initial_cash

            if self.task == "easy":
                score = min(max(profit / (0.01 * self.initial_cash), 0), 1)

            elif self.task == "medium":
                score = min(max(profit / (0.05 * self.initial_cash), 0), 1)

            elif self.task == "hard":
                score = min(max((profit / self.initial_cash) / 0.1, 0), 1)

            reward = score
            reward=float(reward)
        
        return FinenvObservation(
            stock=self.stock,
            market=self.market.value,
            shares=self.shares,
            cash=self.cash,
            price=self.stock_price,
            done=done,
            reward=reward_message(reward),
            metadata={
                "step": self._state.step_count,
                "portfolio_value": new_value
            },
        )

    @property
    def state(self) -> State:
        return self._state