# 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.

"""Supergames environment implementation."""

import copy
from typing import List
from uuid import uuid4

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

try:
    from ..models import SupergamesAction, SupergamesObservation, WorkItem
except ImportError:
    import sys, os
    sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
    from models import SupergamesAction, SupergamesObservation, WorkItem

try:
    from ..tasks import TASKS
    from ..simulator import simulateSprint
except ImportError:
    from tasks import TASKS
    from simulator import simulateSprint


class SupergamesEnvironment(Environment):
    """Environment for staffing and sprint allocation across Supergames tasks."""

    SUPPORTS_CONCURRENT_SESSIONS: bool = True
    MIN_REWARD: float = -100.0
    MAX_REWARD: float = 100.0

    def __init__(self):
        self.stateData = State(episode_id=str(uuid4()), step_count=0)
        # task data
        self.taskId = 1
        self.games = []
        self.workQueue = []
        self.staffPool = None
        self.totalSteps = 0
        self.goal = ""
        # tracking
        self.completedItems: List[WorkItem] = []
        self.cumulativeRevenue = 0.0
        self.crisisResolved = False
        self.done = True
        self.initialWorkQueue = []
        self.initialStaffPool = None
        self.estimatedOptimalRevenue = 1.0
        self.pendingCrisisItems = []

    def buildObservation(self, reward: float | None = None) -> SupergamesObservation:
        return SupergamesObservation(
            taskID=self.taskId,
            currentStep=self.stateData.step_count,
            totalSteps=self.totalSteps,
            games=copy.deepcopy(self.games),
            workQueue=copy.deepcopy(self.workQueue),
            staffPool=copy.deepcopy(self.staffPool),
            crisis=any(item.crisis for item in self.workQueue),            
            goal=self.goal,
            done=self.done,
            reward=reward,
            metadata={
                "completedItems": list(self.completedItems),
                "cumulativeRevenue": round(self.cumulativeRevenue, 2),
            },
        )

    @staticmethod
    def _formatRevenue(amount: float) -> str:
        if abs(amount) >= 1_000_000:
            return f"${amount / 1_000_000:.2f}M"
        if abs(amount) >= 1_000:
            return f"${amount / 1_000:.0f}k"
        return f"${amount:.2f}"

    def _buildSprintSummary(
        self,
        sprintRevenue: float,
        completedItems: List[WorkItem],
    ) -> str:
        unresolvedCriticalBugs = sum(
            1
            for item in self.workQueue
            if item.workType.value == "bug" and int(item.severity) >= 4
        )
        itemLabel = "item" if len(completedItems) == 1 else "items"
        bugLabel = (
            "critical/blocker bug"
            if unresolvedCriticalBugs == 1
            else "critical/blocker bugs"
        )
        summary = (
            f"Sprint {self.stateData.step_count}: "
            f"earned {self._formatRevenue(sprintRevenue)}, "
            f"completed {len(completedItems)} {itemLabel}, "
            f"{unresolvedCriticalBugs} {bugLabel} unresolved"
        )

        maxChurnMult = max((game.churnMult for game in self.games), default=1.0)
        if maxChurnMult > 1.0:
            summary += f" (churn x{maxChurnMult:.2f})"

        return summary

    def _clampReward(self, reward: float) -> float:
        return round(min(self.MAX_REWARD, max(self.MIN_REWARD, reward)), 2)

    @staticmethod
    def _delayImpactPoints(item: WorkItem) -> float:
        if item.impactDelay == 0:
            return 45.0
        if item.impactDelay == 1:
            return 25.0
        return 10.0

    @staticmethod
    def _completedItemPoints(item: WorkItem) -> float:
        severityPoints = int(item.severity) * 3.0
        revenuePoints = min(20.0, item.revenueImpact / 20.0)
        churnPoints = item.churnReduction * 40.0
        return (
            SupergamesEnvironment._delayImpactPoints(item)
            + severityPoints
            + revenuePoints
            + churnPoints
        )

    def computeReward(
        self,
        sprintRevenue: float,
        completedItems: List[WorkItem] | None = None,
        message: str = "ok",
    ) -> float:
        completedItems = completedItems or []

        if message.startswith("Overallocation"):
            return -70.0
        if message.startswith("Unknown work item"):
            return -60.0

        reward = min(35.0, sprintRevenue / 100_000.0)
        reward += sum(self._completedItemPoints(item) for item in completedItems)

        unresolvedPenalty = 0.0
        for item in self.workQueue:
            if item.workType.value != "bug":
                continue
            if int(item.severity) == 5:
                unresolvedPenalty += 18.0
            elif int(item.severity) == 4:
                unresolvedPenalty += 10.0
        reward -= unresolvedPenalty

        churnPenalty = sum(max(0.0, game.churnMult - 1.0) * 40.0 for game in self.games)
        reward -= churnPenalty

        if self.done:
            if not any(item.workType.value == "bug" and int(item.severity) >= 4 for item in self.workQueue):
                reward += 15.0
            if self.taskId == 4:
                reward += 20.0 if self.crisisResolved else -40.0

        return self._clampReward(reward)

    def reset(self, task_id: int = 1, seed: int = 42) -> SupergamesObservation:
        if task_id not in TASKS:
            raise ValueError(f"Unknown task_id: {task_id}")

        self.stateData = State(episode_id=str(uuid4()), step_count=0)
        self.taskId = task_id

        games, workQueue, staffPool, totalSteps, goal = TASKS[task_id]["generate"](seed)
        self.games = copy.deepcopy(games)
        self.workQueue = copy.deepcopy(workQueue)
        self.staffPool = copy.deepcopy(staffPool)
        self.totalSteps = totalSteps
        self.goal = goal

        # For task 4, crisis work should only become available starting sprint 2.
        self.pendingCrisisItems = []
        if self.taskId == 4:
            self.pendingCrisisItems = [item for item in self.workQueue if item.crisis]
            self.workQueue = [item for item in self.workQueue if not item.crisis]

        self.completedItems = []
        self.cumulativeRevenue = 0.0
        self.crisisResolved = False
        self.done = False
        self.initialWorkQueue = copy.deepcopy(self.workQueue)
        self.initialStaffPool = copy.deepcopy(self.staffPool)

        # Coarse upper bound for multi-sprint normalization.
        baseRevenue = sum(game.monthlyRevenue for game in self.games) * self.totalSteps
        impactBound = 0.0
        for item in self.workQueue + self.pendingCrisisItems:
            activeSprints = max(0, self.totalSteps - item.impactDelay)
            impactBound += item.revenueImpact * (activeSprints / max(1, self.totalSteps))
        self.estimatedOptimalRevenue = max(1.0, round(baseRevenue + impactBound, 2))

        return self.buildObservation(reward=0.0)

    def step(self, action: SupergamesAction) -> SupergamesObservation:
        if self.done:
            return self.buildObservation(reward=0.0)

        self.stateData.step_count += 1
        reasoning = action.reasoning.strip()

        if self.taskId == 4 and self.stateData.step_count >= 2 and self.pendingCrisisItems:
            self.workQueue.extend(self.pendingCrisisItems)
            self.pendingCrisisItems = []

        (
            self.games,
            self.workQueue,
            completedItems,
            sprintRevenue,
            message,
        ) = simulateSprint(
            self.games,
            self.workQueue,
            self.staffPool,
            action,
            self.stateData.step_count,
        )

        self.cumulativeRevenue += sprintRevenue
        self.completedItems.extend(completedItems)
        sprintSummary = self._buildSprintSummary(sprintRevenue, completedItems)

        if self.taskId == 4 and any(item.crisis for item in completedItems):
            self.crisisResolved = True

        self.done = self.stateData.step_count >= self.totalSteps
        reward = self.computeReward(sprintRevenue, completedItems, message)

        obs = self.buildObservation(reward=reward)
        obs.metadata.update(
            {
                "stepRevenue": round(sprintRevenue, 2),
                "message": message,
                "sprint_summary": sprintSummary,
            }
        )
        if reasoning:
            obs.metadata["agent_reasoning"] = reasoning
        return obs

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