from __future__ import annotations

import traceback
from typing import Any, Dict

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

from cloudenv.models import CloudEnvAction, CloudEnvObservation, CloudEnvState


class ClientEnvironment(
    EnvClient[CloudEnvAction, CloudEnvObservation, CloudEnvState]
):
    # ---------------------------
    # INIT ACTION
    # ---------------------------

    def initialize_environment(
        self, *, task_id: str, episode_id: str
    ) -> StepResult[CloudEnvObservation]:
        try:
            action = CloudEnvAction(
                task_id=task_id,
                episode_id=episode_id,
                service="internal",  # safer than enum
                operation="InitializeEnvironment",
                payload={},
            )
            return self.step(action)

        except Exception:
            print("[CLIENT ERROR] initialize_environment failed", flush=True)
            traceback.print_exc()

            return StepResult(
                observation=CloudEnvObservation(
                    episode_id=episode_id,
                    task_id=task_id,
                    step_count=0,
                    operation_result={"error": "init failed"},
                    reward=0.0,
                    done=True,
                    metadata={},
                    current_pipeline_state={},
                ),
                reward=0.0,
                done=True,
            )

    # ---------------------------
    # SERIALIZATION
    # ---------------------------

    def _step_payload(self, action: CloudEnvAction) -> Dict[str, Any]:
        try:
            return action.model_dump(mode="json", exclude_none=True)
        except Exception:
            print(
                "[CLIENT ERROR] step payload serialization failed", flush=True
            )
            traceback.print_exc()
            return {}

    # ---------------------------
    # PARSE RESULT
    # ---------------------------

    def _parse_result(
        self, payload: Dict[str, Any]
    ) -> StepResult[CloudEnvObservation]:
        try:
            payload = payload or {}
            reward = payload.get("reward", 0.0)

            observation = CloudEnvObservation(
                episode_id=payload.get("episode_id", ""),
                task_id=payload.get("task_id", ""),
                step_count=payload.get("step_count", 0),
                operation_result=payload.get("operation_result", {}),
                reward=float(reward) if reward is not None else 0.0,
                done=bool(payload.get("done", False)),
                metadata=payload.get("metadata", {}),
                current_pipeline_state=payload.get(
                    "current_pipeline_state", {}
                ),
            )

            return StepResult(
                observation=observation,
                reward=observation.reward,
                done=observation.done,
            )

        except Exception:
            print("[CLIENT ERROR] parse_result failed", flush=True)
            traceback.print_exc()

            return StepResult(
                observation=CloudEnvObservation(
                    episode_id="",
                    task_id="",
                    step_count=0,
                    operation_result={"error": "parse failure"},
                    reward=0.0,
                    done=True,
                    metadata={},
                    current_pipeline_state={},
                ),
                reward=0.0,
                done=True,
            )

    # ---------------------------
    # PARSE STATE
    # ---------------------------

    def _parse_state(self, payload: Dict[str, Any]) -> CloudEnvState:
        try:
            payload = payload or {}

            return CloudEnvState(
                task_id=payload.get("task_id"),
                episode_id=payload.get("episode_id"),
                step_count=payload.get("step_count", 0),
                reward_history=payload.get("reward_history", []),
            )

        except Exception:
            print("[CLIENT ERROR] parse_state failed", flush=True)
            traceback.print_exc()

            return CloudEnvState(
                task_id=None,
                episode_id=None,
                step_count=0,
                reward_history=[],
            )