import os

from agentscope.agent import ReActAgent
from agentscope.model import DashScopeChatModel
from agentscope.formatter import DashScopeChatFormatter
from agentscope.tool import Toolkit, execute_python_code
from agentscope.pipeline import stream_printing_messages

from agentscope_runtime.engine import AgentApp
from agentscope_runtime.engine.schemas.agent_schemas import AgentRequest
from agentscope_runtime.adapters.agentscope.memory import (
    AgentScopeSessionHistoryMemory,
)
from agentscope_runtime.engine.services.agent_state import (
    InMemoryStateService,
)
from agentscope_runtime.engine.services.session_history import (
    InMemorySessionHistoryService,
)

agent_app = AgentApp(
    app_name="Friday",
    app_description="A helpful assistant",
)


@agent_app.init
async def init_func(self):
    self.state_service = InMemoryStateService()
    self.session_service = InMemorySessionHistoryService()

    await self.state_service.start()
    await self.session_service.start()


@agent_app.shutdown
async def shutdown_func(self):
    await self.state_service.stop()
    await self.session_service.stop()


@agent_app.query(framework="agentscope")
async def query_func(
    self,
    msgs,
    request: AgentRequest = None,
    **kwargs,
):
    session_id = request.session_id
    user_id = request.user_id

    state = await self.state_service.export_state(
        session_id=session_id,
        user_id=user_id,
    )

    toolkit = Toolkit()
    toolkit.register_tool_function(execute_python_code)

    agent = ReActAgent(
        name="Friday",
        model=DashScopeChatModel(
            "qwen-turbo",
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            stream=True,
        ),
        sys_prompt="You're a helpful assistant named Friday.",
        toolkit=toolkit,
        memory=AgentScopeSessionHistoryMemory(
            service=self.session_service,
            session_id=session_id,
            user_id=user_id,
        ),
        formatter=DashScopeChatFormatter(),
    )
    agent.set_console_output_enabled(enabled=False)

    if state:
        agent.load_state_dict(state)

    async for msg, last in stream_printing_messages(
        agents=[agent],
        coroutine_task=agent(msgs),
    ):
        yield msg, last

    state = agent.state_dict()

    await self.state_service.save_state(
        user_id=user_id,
        session_id=session_id,
        state=state,
    )


agent_app.run(host="127.0.0.1", port=8090)

from agentscope_runtime.sandbox import BaseSandbox

with BaseSandbox() as box:
    # By default, pulls `agentscope/runtime-sandbox-base:latest` from DockerHub
    print(box.list_tools()) # List all available tools
    print(box.run_ipython_cell(code="print('hi')"))  # Run Python code
    print(box.run_shell_command(command="echo hello"))  # Run shell command
    input("Press Enter to continue...")

from agentscope_runtime.sandbox import GuiSandbox

with GuiSandbox() as box:
    # By default, pulls `agentscope/runtime-sandbox-gui:latest` from DockerHub
    print(box.list_tools()) # List all available tools
    print(box.desktop_url)  # Web desktop access URL
    print(box.computer_use(action="get_cursor_position"))  # Get mouse cursor position
    print(box.computer_use(action="get_screenshot"))       # Capture screenshot
    input("Press Enter to continue...")

from agentscope_runtime.sandbox import BrowserSandbox

with BrowserSandbox() as box:
    # By default, pulls `agentscope/runtime-sandbox-browser:latest` from DockerHub
    print(box.list_tools()) # List all available tools
    print(box.desktop_url)  # Web desktop access URL
    box.browser_navigate("https://www.google.com/")  # Open a webpage
    input("Press Enter to continue...")

from agentscope_runtime.sandbox import FilesystemSandbox

with FilesystemSandbox() as box:
    # By default, pulls `agentscope/runtime-sandbox-filesystem:latest` from DockerHub
    print(box.list_tools()) # List all available tools
    print(box.desktop_url)  # Web desktop access URL
    box.create_directory("test")  # Create a directory
    input("Press Enter to continue...")

from agentscope_runtime.sandbox import MobileSandbox

with MobileSandbox() as box:
    # By default, pulls 'agentscope/runtime-sandbox-mobile:latest' from DockerHub
    print(box.list_tools()) # List all available tools
    print(box.mobile_get_screen_resolution()) # Get the screen resolution
    print(box.mobile_tap([500, 1000])) # Tap at coordinate (500, 1000)
    print(box.mobile_input_text("Hello from AgentScope!")) # Input text
    print(box.mobile_key_event(3)) # Sends a HOME key event (KeyCode: 3)
    screenshot_result = box.mobile_get_screenshot() # Get the current screenshot
    input("Press Enter to continue...")

from agentscope_runtime.engine.deployers import LocalDeployManager

# Create deployment manager
deployer = LocalDeployManager(
    host="0.0.0.0",
    port=8090,
)

# Deploy the app as a streaming service
deploy_result = await app.deploy(deployer=deployer)

from openai import OpenAI

client = OpenAI(base_url="http://0.0.0.0:8090/compatible-mode/v1")

response = client.responses.create(
  model="any_name",
  input="What is the weather in Beijing?"
)

print(response)

from agentscope_runtime.engine.deployers import ModelStudioDeployManager
# Create deployment manager
deployer = ModelstudioDeployManager(
    oss_config=OSSConfig(
        access_key_id=os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID"),
        access_key_secret=os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET"),
    ),
    modelstudio_config=ModelstudioConfig(
        workspace_id=os.environ.get("MODELSTUDIO_WORKSPACE_ID"),
        access_key_id=os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID"),
        access_key_secret=os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET"),
        dashscope_api_key=os.environ.get("DASHSCOPE_API_KEY"),
    ),
)

# Deploy to ModelStudio
result = await app.deploy(
    deployer,
    deploy_name="agent-app-example",
    telemetry_enabled=True,
    requirements=["agentscope", "fastapi", "uvicorn"],
    environment={
        "PYTHONPATH": "/app",
        "DASHSCOPE_API_KEY": os.environ.get("DASHSCOPE_API_KEY"),
    },
)