import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from engine import (
    GameState, INITIAL_DEMAND, SALE_PRICE, SHIPPING_COST,
    CONTAINER_SIZE, ALL_WAREHOUSE_IDS
)
import json
import copy


class ComponentRubric:
    def evaluate(self, state, prev_state):
        return 0.0


class FinancialRubric(ComponentRubric):
    def evaluate(self, state, prev_state):
        # Reward profit scalability
        diff = state.balance - prev_state.balance
        return diff / 1000.0


class StockoutRubric(ComponentRubric):
    def evaluate(self, state, prev_state):
        # Hard to game: heavy penalty for failing the core objective (serving demand)
        penalty = 0
        for region, demand in state.demands.items():
            if state.inventory[region] < demand:
                missed = demand - state.inventory[region]
                penalty -= (missed * 10)  # Heavy negative scaling
        return penalty


class OverstockRubric(ComponentRubric):
    def evaluate(self, state, prev_state):
        # Prevents agent from gaming the stockout penalty by just buying infinite inventory
        penalty = 0
        for region, inventory in state.inventory.items():
            if inventory > state.demands[region] * 2:
                excess = inventory - (state.demands[region] * 2)
                penalty -= (excess * 2)
        return penalty


class SupplyChainEnv:
    """
    OpenEnv-compliant environment wrapper for the supply chain simulation.
    Follows the standard Gym-style API (reset, step, state).
    Features a composable Rubric reward system.
    """

    def __init__(self, max_turns=24):
        self._state = GameState()
        self.max_turns = max_turns
        self.rubrics = [FinancialRubric(), StockoutRubric(), OverstockRubric()]

    @property
    def state(self):
        """Return the current observation (Gym-style API compliance)."""
        return self._get_obs()

    def reset(self, seed=None, options=None):
        self._state.reset()
        return self._get_obs(), {"balance": self._state.balance}

    def _get_obs(self):
        obs = {
            "Month": self._state.turn,
            "Funds": self._state.balance,
            "Demands": self._state.demands,
            "Inventory": self._state.inventory,
            "Factory_Costs": {f["id"]: f["currentCost"] for f in self._state.factoryCosts},
            "Port_Delays": {p["id"]: ("CLOSED" if p["isClosed"] else p["currentDelay"]) for p in self._state.portDelays},
        }
        return json.dumps(obs)

    def _compute_reward(self, prev_state):
        total_reward = 0
        info_rubrics = {}
        for rubric in self.rubrics:
            r = rubric.evaluate(self._state, prev_state)
            total_reward += r
            info_rubrics[rubric.__class__.__name__] = r
        return total_reward, info_rubrics

    def step(self, action: str):
        done = False
        info = {}

        # Deepcopy state to compare for rubrics
        prev_state = copy.deepcopy(self._state)

        try:
            act_dict = json.loads(action)
            action_type = act_dict.get("action", "advance")

            if action_type == "order":
                factory_id = act_dict.get("factory")
                quantity = act_dict.get("quantity", 0)
                port_id = act_dict.get("port")
                dest = act_dict.get("dest")

                # Manual order logic equivalent
                factory = next((f for f in self._state.factoryCosts if f["id"] == factory_id), None)
                port = next((p for p in self._state.portDelays if p["id"] == port_id), None)
                if factory and port and not port["isClosed"]:
                    cost = factory["currentCost"] * quantity
                    shipping = ((quantity + CONTAINER_SIZE - 1) // CONTAINER_SIZE) * SHIPPING_COST
                    if self._state.balance >= cost + shipping:
                        self._state.balance -= (cost + shipping)
                        self._state.month_goods_costs += cost
                        self._state.month_shipping_costs += shipping
                        self._state.inTransit.append({
                            "factory": factory_id, "port": port_id, "dest": dest,
                            "quantity": quantity, "delayLeft": port["currentDelay"]
                        })
            elif action_type == "advance":
                # Replicate advance turn
                # 1. Shipments arrive
                new_transit = []
                for s in self._state.inTransit:
                    if s["delayLeft"] <= 1:
                        self._state.inventory[s["dest"]] += s["quantity"]
                    else:
                        s["delayLeft"] -= 1
                        new_transit.append(s)
                self._state.inTransit = new_transit

                # 2. Fulfill demand for ALL warehouse regions (fixed from 3-region bug)
                for r in ALL_WAREHOUSE_IDS:
                    demand = self._state.demands.get(r, 0)
                    sold = min(demand, self._state.inventory[r])
                    self._state.inventory[r] -= sold
                    self._state.balance += sold * SALE_PRICE

                self._state.turn += 1

        except Exception as e:
            info["error"] = str(e)
            return self._get_obs(), -50.0, False, False, info  # Parse penalty

        reward, rubrics = self._compute_reward(prev_state)
        info["rubrics"] = rubrics

        if self._state.turn >= self.max_turns or self._state.balance <= 0:
            done = True

        return self._get_obs(), reward, done, False, info