# Copyright 2025 Dimensional Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import os
from dimos.robot.robot import MockRobot
import tests.test_header

import time
from dotenv import load_dotenv
from dimos.agents.cerebras_agent import CerebrasAgent
from dimos.robot.unitree.unitree_skills import MyUnitreeSkills
from dimos.web.robot_web_interface import RobotWebInterface
from dimos.skills.kill_skill import KillSkill
from dimos.skills.navigation import NavigateWithText, GetPose, NavigateToGoal
from dimos.skills.visual_navigation_skills import FollowHuman
import reactivex as rx
import reactivex.operators as ops
from dimos.stream.audio.pipelines import tts, stt
from dimos.web.websocket_vis.server import WebsocketVis
import threading
from dimos.types.vector import Vector
from dimos.skills.speak import Speak
from dimos.skills.skills import SkillLibrary

# Load API key from environment
load_dotenv()

robot = MockRobot()

# Create a subject for agent responses
agent_response_subject = rx.subject.Subject()
agent_response_stream = agent_response_subject.pipe(ops.share())

text_streams = {
    "agent_responses": agent_response_stream,
}

web_interface = RobotWebInterface(
    port=5556,
    text_streams=text_streams,
)

robot_skills = SkillLibrary()

# Create a CerebrasAgent instance
agent = CerebrasAgent(
    dev_name="test_cerebras_agent",
    # input_query_stream=stt_node.emit_text(),
    input_query_stream=web_interface.query_stream,
    skills=robot_skills,
    system_query="""You are an agent controlling a virtual robot. When given a query, respond by using the appropriate tool calls if needed to execute commands on the robot.

IMPORTANT INSTRUCTIONS:
1. Each tool call must include the exact function name and appropriate parameters
2. If a function needs parameters like 'distance' or 'angle', be sure to include them
3. If you're unsure which tool to use, choose the most appropriate one based on the user's query
4. Parse the user's instructions carefully to determine correct parameter values

When you need to call a skill or tool, ALWAYS respond ONLY with a JSON object in this exact format: {"name": "SkillName", "arguments": {"arg1": "value1", "arg2": "value2"}}

Example: If the user asks to spin right by 90 degrees, output ONLY the following: {"name": "SpinRight", "arguments": {"degrees": 90}}""",
    model_name="llama-4-scout-17b-16e-instruct",
)

robot_skills.add(KillSkill)
robot_skills.add(NavigateWithText)
robot_skills.add(FollowHuman)
robot_skills.add(GetPose)
robot_skills.add(NavigateToGoal)
robot_skills.create_instance("KillSkill", robot=robot, skill_library=robot_skills)
robot_skills.create_instance("NavigateWithText", robot=robot)
robot_skills.create_instance("FollowHuman", robot=robot)
robot_skills.create_instance("GetPose", robot=robot)
robot_skills.create_instance("NavigateToGoal", robot=robot)


# Subscribe to agent responses and send them to the subject
agent.get_response_observable().subscribe(lambda x: agent_response_subject.on_next(x))

# print(f"Registered skills: {', '.join([skill.__name__ for skill in robot_skills.skills])}")
print("Cerebras agent demo initialized. You can now interact with the agent via the web interface.")

web_interface.run()
