from __future__ import annotations

from typing import Dict

from openenv.core import EnvClient
from openenv.core.client_types import StepResult

from models import (
    ActionType,
    B2BSupportPayload,
    B2BSupportTriageAction,
    B2BSupportTriageObservation,
    B2BSupportTriageState,
    RewardBreakdown,
    VisibleTicket,
)


class B2BSupportTriageEnv(
    EnvClient[B2BSupportTriageAction, B2BSupportTriageObservation, B2BSupportTriageState]
):
    def _step_payload(self, action: B2BSupportTriageAction) -> Dict:
        return {
            "action_type": action.action_type.value,
            "ticket_id": action.ticket_id,
            "payload": action.payload.model_dump(exclude_none=True),
            "metadata": action.metadata,
        }

    def _parse_result(self, payload: Dict) -> StepResult[B2BSupportTriageObservation]:
        obs_data = payload.get("observation", {})
        reward_breakdown = obs_data.get("reward_breakdown", {})

        observation = B2BSupportTriageObservation(
            task_id=obs_data.get("task_id", "easy"),
            step_index=obs_data.get("step_index", 0),
            max_steps=obs_data.get("max_steps", 0),
            visible_ticket=VisibleTicket(**obs_data.get("visible_ticket", {})),
            current_plan=obs_data.get("current_plan", []),
            applied_decisions=obs_data.get("applied_decisions", {}),
            last_action_error=obs_data.get("last_action_error"),
            progress_score=obs_data.get("progress_score", 0.0),
            reward_breakdown=RewardBreakdown(**reward_breakdown),
            done=payload.get("done", False),
            reward=payload.get("reward"),
            metadata=obs_data.get("metadata", {}),
        )

        return StepResult(
            observation=observation,
            reward=payload.get("reward"),
            done=payload.get("done", False),
        )

    def _parse_state(self, payload: Dict) -> B2BSupportTriageState:
        return B2BSupportTriageState(
            episode_id=payload.get("episode_id", ""),
            step_count=payload.get("step_count", 0),
            task_id=payload.get("task_id", ""),
            seed=payload.get("seed"),
            max_steps=payload.get("max_steps", 0),
            cumulative_reward=payload.get("cumulative_reward", 0.0),
            last_action_error=payload.get("last_action_error"),
            applied_decisions=payload.get("applied_decisions", {}),
            action_history=payload.get("action_history", []),
            completion_flags=payload.get("completion_flags", {}),
        )


__all__ = [
    "ActionType",
    "B2BSupportPayload",
    "B2BSupportTriageAction",
    "B2BSupportTriageEnv",
    "B2BSupportTriageObservation",
    "B2BSupportTriageState",
]