import time

from linksoul_agentsdk import (
    AgentAuthCallback,
    AgentParam,
    AgentSdk,
    ExtSkillInvokeCallback,
    ExtSkillQueryCallback,
    ExtSkillRequest,
    IdGenerator,
)

url        = "wss://open.agibot.com/api/V1/open-portal/app/wss/agent-sdk"
app_id     = "<appId from the LinkSoul open platform>"
app_key    = "<appKey from the LinkSoul open platform>"
app_secret = "<appSecret from the LinkSoul open platform>"
agent_id   = "<agentId from your custom agent on the LinkSoul platform>"

# 1. Create the SDK instance
agent_sdk = AgentSdk.create(url, app_id, app_key, app_secret)


# 2. Register auth callback
class MyAuthCallback(AgentAuthCallback):
    def on_auth_state(self, app_id: str, code: int, msg: str) -> None:
        print(f"on_auth_state => app_id: {app_id}, code: {code}, msg: {msg}")


agent_sdk.register_auth(MyAuthCallback())

# 3. Open the WebSocket connection
agent_sdk.initialize()
time.sleep(2)

# 4. Register an extended-skill request instance
request_id = IdGenerator.generate_ext_skill_request_id()
ext_skill_request = ExtSkillRequest(agent_sdk, agent_id, request_id)
agent_sdk.register_ext_skill(ext_skill_request)


# 4.1 Query the available extended-skill list
class MyQueryCallback(ExtSkillQueryCallback):
    def on_query_result(self, agent_id: str, request_id: str, event_id: str,
                        code: int, msg: str, result) -> None:
        print(f"on_query_result => agent_id={agent_id}, request_id={request_id}, "
              f"event_id={event_id}, code={code}, msg={msg}, result={result}")


query_event_id = ext_skill_request.query(MyQueryCallback(), 2000)
print(f"query event_id: {query_event_id}")


# 4.2 Trigger a skill to run
ext_skill_id = "test_ext_skill_id_001"           # extended skill ID
input_param = (                                   # skill input
    AgentParam.create()
    .set_string("role", "male")
    .set_integer("threshold", 0)
)


class MyInvokeCallback(ExtSkillInvokeCallback):
    def on_invoke_ack(self, agent_id: str, request_id: str, event_id: str,
                      code: int, msg: str) -> None:
        # stage 1: gateway accepted
        print(f"on_invoke_ack => agent_id={agent_id}, request_id={request_id}, "
              f"event_id={event_id}, code={code}, msg={msg}")

    def on_state(self, agent_id: str, request_id: str, event_id: str,
                 state_name, state_value) -> None:
        # stage 2: execution-time state report (zero or more times)
        print(f"on_state => agent_id={agent_id}, request_id={request_id}, "
              f"event_id={event_id}, state_name={state_name}, state_value={state_value}")

    def on_invoke_result(self, agent_id: str, request_id: str, event_id: str,
                         code: int, msg: str, param) -> None:
        # stage 3: skill execution output
        print(f"on_invoke_result => agent_id={agent_id}, request_id={request_id}, "
              f"event_id={event_id}, code={code}, msg={msg}, param={param}")


invoke_event_id = ext_skill_request.invoke(ext_skill_id, input_param, MyInvokeCallback(), 2000)
print(f"invoke event_id: {invoke_event_id}")

time.sleep(5)

# 5. Unregister to release
agent_sdk.unregister_ext_skill(request_id)
# agent_sdk.release()