from __future__ import annotations

from typing import Any

from pydantic import Field
from openenv.core.env_server import Action, Environment, Observation, State

from agents import SmartCareerAgent
from logging_utils import get_logger

logger = get_logger("environment")


class CareerAction(Action):
    query: str = ""


class CareerObservation(Observation):
    output: dict[str, Any]


class CareerState(State):
    last_query: str = ""
    last_output: dict[str, Any] = Field(default_factory=dict)


class CareerEnvironment(Environment[CareerAction, CareerObservation, CareerState]):
    def __init__(self) -> None:
        super().__init__()
        self._agent = SmartCareerAgent()
        self._state = CareerState()
        logger.info("CareerEnvironment initialized.")

    def reset(
        self,
        seed: int | None = None,
        episode_id: str | None = None,
        **kwargs: Any,
    ) -> CareerObservation:
        logger.info("reset called seed=%r episode_id=%r", seed, episode_id)
        self._agent = SmartCareerAgent()
        self._state = CareerState(episode_id=episode_id)
        return CareerObservation(output={"status": "ok"}, done=False, reward=None)

    def step(
        self,
        action: CareerAction,
        timeout_s: float | None = None,
        **kwargs: Any,
    ) -> CareerObservation:
        query = action.query or ""
        logger.info("step called query=%r timeout_s=%r", query, timeout_s)
        result = self._agent.respond(query)

        self._state.step_count += 1
        self._state.last_query = query
        self._state.last_output = result

        logger.info("step completed step_count=%s intent=%s", self._state.step_count, result.get("intent"))
        return CareerObservation(output=result, done=False, reward=None)

    @property
    def state(self) -> CareerState:
        logger.info("state requested step_count=%s", self._state.step_count)
        return self._state