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

"""
Data models for the Fbot Agent Sim Environment.

The fbot_agent_sim environment is a simple test environment that echoes back messages.
"""

from openenv.core.env_server.types import Action, Observation
from pydantic import Field


from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, field
import random
import math

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


# ---------------------------------------------------------------------
# Data structures for the simulated world
# ---------------------------------------------------------------------
@dataclass
class Pose:
    """Simple 2D pose (x, y, theta in radians)."""
    x: float = 0.0
    y: float = 0.0
    theta: float = 0.0

    def distance_to(self, other: "Pose") -> float:
        return math.hypot(self.x - other.x, self.y - other.y)

    def __str__(self) -> str:
        return f"Pose(x={self.x:.2f}, y={self.y:.2f}, theta={self.theta:.2f})"


@dataclass
class Object:
    """Represents an object in the world."""
    name: str
    pose: Pose
    is_pickable: bool = True
    description: str = ""


@dataclass
class Person:
    """Represents a person in the world."""
    name: str
    appearance: str          # physical description (e.g., "tall man with red shirt")
    face_uuid: Optional[str] = None   # for face recognition
    pose: Pose = field(default_factory=Pose)


class RobotState:
    """Current state of the simulated robot."""
    def __init__(self):
        self.pose = Pose(0.0, 0.0, 0.0)
        self.held_object: Optional[str] = None
        self.emotion: str = "neutral"
        self.last_speech: str = ""
        self.heard_speech: Optional[str] = None  # from listen_something


# ---------------------------------------------------------------------
# State – full world state (not directly observed by the agent)
# ---------------------------------------------------------------------
@dataclass
class WorldState(State):
    """Complete simulation state."""
    robot: RobotState
    objects: Dict[str, Object]
    people: Dict[str, Person]
    named_locations: Dict[str, Pose]   # e.g., "kitchen", "living_room"


class FbotAgentSimAction(Action):
    """A discriminated union for all robot actions (tools)."""
    tool: str
    # Navigation
    location_name: Optional[str] = None          # for navigate_to_pose, query_pose_by_name
    distance: Optional[float] = None             # for move_forward
    angle: Optional[float] = None                # for rotate
    object_name: Optional[str] = None            # for detect_and_approach_object, detect_and_pick_object, etc.
    object_class: Optional[str] = None
    person_description: Optional[str] = None     # for detect_and_approach_person_by_description
    person_name: Optional[str] = None            # for detect_and_approach_unknown_person_by_name
    # Vision
    question: Optional[str] = None               # for analyze_scene
    count_object: Optional[str] = None           # for count_objects
    count_person_desc: Optional[str] = None      # for count_people
    search_name: Optional[str] = None            # for search_person
    face_uuid: Optional[str] = None              # for find_person_saved_by_face, save_person_face
    # Manipulation
    location_pose: Optional[str] = None          # for place_object (named location or "table", etc.)
    target_user: Optional[str] = None            # for give_object_to_user (person name)
    # Utility
    emotion: Optional[str] = None                # for set_emotion
    text: Optional[str] = None                   # for say_something, get_question_answer
    # Transform
    target_frame: Optional[str] = None           # for transform_pose
    source_pose: Optional[Pose] = None


class FbotAgentSimObservation(Observation):
    """Observation returned after each action."""
    robot_pose: Pose
    held_object: Optional[str]
    near_objects: List[str]          # names of objects within 2m
    near_people: List[str]           # names of people within 5m
    last_speech: str                  # what the robot just said
    heard_speech: Optional[str]       # what the robot heard (listen)
    emotion: str
    success: bool
    message: str                      # textual description of outcome