"""主动交互——任务流（对齐 Java ``TaskFlowRequest``）。

调用流程：``start_request`` 必须 → 任意次 ``task_request`` 可选 → ``end_request`` 可选；
任何阶段都可以 ``interrupt_request``。
"""

from __future__ import annotations

import json
import logging
import time
from typing import TYPE_CHECKING, Any, Dict, Optional

from ..enums import AgentEventType, AgentPolicy, FlowMode
from ..id_generator import IdGenerator

if TYPE_CHECKING:
    from ..agent_sdk import AgentSdk
    from ..client.linksky_client import LinkskyClient


_log = logging.getLogger("linksoul_agentsdk.active")


_MIN_TIMEOUT_MS = 50
_CONN_CLOSED_CODE = 1000
_CONN_CLOSED_MSG = "Agentsdk connection is closed"
_BAD_JSON_CODE = 1001
_BAD_JSON_MSG = "payload is not valid JSON"


# ---------------------------------------------------------------------------
# Flow callbacks
# ---------------------------------------------------------------------------


class _FlowCallbackBase:
    """所有 Flow 回调的共享行为：携带 ``timeout``。"""

    def __init__(self) -> None:
        self.timeout: int = _MIN_TIMEOUT_MS

    def set_timeout(self, timeout: int) -> None:
        self.timeout = max(timeout, _MIN_TIMEOUT_MS)


class FlowStartCallback(_FlowCallbackBase):
    def on_start_request_ack(self, flow_id: str, event_id: str, code: int, msg: str) -> None:
        raise NotImplementedError

    def on_start_execute_ack(self, flow_id: str, event_id: str, code: int, msg: str) -> None:
        raise NotImplementedError


class FlowTaskCallback(_FlowCallbackBase):
    def on_task_request_ack(self, flow_id: str, event_id: str, code: int, msg: str) -> None:
        raise NotImplementedError

    def on_task_execute_ack(self, flow_id: str, event_id: str, code: int, msg: str) -> None:
        raise NotImplementedError


class FlowEndCallback(_FlowCallbackBase):
    def on_end_request_ack(self, flow_id: str, event_id: str, code: int, msg: str) -> None:
        raise NotImplementedError


class FlowInterruptCallback(_FlowCallbackBase):
    def on_interrupt_request_ack(self, flow_id: str, event_id: str, code: int, msg: str) -> None:
        raise NotImplementedError


# ---------------------------------------------------------------------------
# TaskFlowRequest
# ---------------------------------------------------------------------------


class TaskFlowRequest:
    """``register_task_flow(...)`` 注册后即可调用 ``start_request`` 等方法。"""

    def __init__(
        self,
        agent_sdk: "AgentSdk",
        agent_id: str,
        flow_id: str,
        flow_mode: FlowMode = FlowMode.COMPLEX,
    ) -> None:
        self._agent_sdk = agent_sdk
        self._agent_id = agent_id
        self._flow_id = flow_id
        self._flow_mode = flow_mode

        # 由 AgentSdk.register_task_flow 注入
        self._linksky_client: Optional["LinkskyClient"] = None

        self._start_callback: Optional[FlowStartCallback] = None
        self._end_callback: Optional[FlowEndCallback] = None
        self._interrupt_callback: Optional[FlowInterruptCallback] = None
        self._event_id_to_task_cb: Dict[str, FlowTaskCallback] = {}
        self._update_ts = int(time.time() * 1000)

    # ----- accessors -----

    @property
    def agent_id(self) -> str:
        return self._agent_id

    @property
    def flow_id(self) -> str:
        return self._flow_id

    @property
    def flow_mode(self) -> FlowMode:
        return self._flow_mode

    @property
    def update_ts(self) -> int:
        return self._update_ts

    def set_linksky_client(self, client: "LinkskyClient") -> None:
        self._linksky_client = client

    # 给外部 Mgr 用：从 event_id 取出/移除 task 回调（一次任务通常对应一个）
    def get_task_callback(self, event_id: str) -> Optional[FlowTaskCallback]:
        return self._event_id_to_task_cb.get(event_id)

    def remove_task_callback(self, event_id: str) -> None:
        self._event_id_to_task_cb.pop(event_id, None)

    @property
    def start_callback(self) -> Optional[FlowStartCallback]:
        return self._start_callback

    @property
    def end_callback(self) -> Optional[FlowEndCallback]:
        return self._end_callback

    @property
    def interrupt_callback(self) -> Optional[FlowInterruptCallback]:
        return self._interrupt_callback

    # ----- request methods -----

    def start_request(
        self,
        payload: Optional[str],
        start_callback: FlowStartCallback,
        timeout: int,
    ) -> str:
        start_callback.set_timeout(timeout)
        self._start_callback = start_callback

        event_id = IdGenerator.generate_event_id()
        envelope: Dict[str, Any] = {
            "type": AgentEventType.AGENTSDK_FLOW_START_REQUEST.value,
            "agentId": self._agent_id,
            "agentMode": AgentPolicy.ACTIVE.value,
            "flowId": self._flow_id,
            "mode": self._flow_mode.value,
            "eventId": event_id,
        }

        if self._flow_mode is FlowMode.SIMPLE and not payload:
            _log.error("start_request => SIMPLE mode payload is empty, flow_id %s", self._flow_id)
            start_callback.on_start_request_ack(self._flow_id, event_id, _CONN_CLOSED_CODE, _CONN_CLOSED_MSG)
            return event_id

        if payload:
            try:
                envelope["payload"] = json.loads(payload)
            except json.JSONDecodeError as exc:
                _log.error("start_request => bad payload, flow_id %s, exc %s", self._flow_id, exc)
                start_callback.on_start_request_ack(self._flow_id, event_id, _BAD_JSON_CODE, _BAD_JSON_MSG)
                return event_id

        text = json.dumps(envelope, ensure_ascii=False)
        self._send(text, event_id, lambda code, msg: start_callback.on_start_request_ack(self._flow_id, event_id, code, msg))
        return event_id

    def task_request(
        self,
        payload: str,
        task_callback: FlowTaskCallback,
        timeout: int,
    ) -> str:
        task_callback.set_timeout(timeout)
        event_id = IdGenerator.generate_event_id()
        self._event_id_to_task_cb[event_id] = task_callback

        envelope: Dict[str, Any] = {
            "type": AgentEventType.AGENTSDK_FLOW_TASK_REQUEST.value,
            "agentId": self._agent_id,
            "agentMode": AgentPolicy.ACTIVE.value,
            "flowId": self._flow_id,
            "eventId": event_id,
        }
        try:
            envelope["payload"] = json.loads(payload)
        except json.JSONDecodeError as exc:
            _log.error("task_request => bad payload, flow_id %s, exc %s", self._flow_id, exc)
            task_callback.on_task_request_ack(self._flow_id, event_id, _BAD_JSON_CODE, _BAD_JSON_MSG)
            return event_id

        text = json.dumps(envelope, ensure_ascii=False)
        self._send(text, event_id, lambda code, msg: task_callback.on_task_request_ack(self._flow_id, event_id, code, msg))
        return event_id

    def end_request(self, end_callback: FlowEndCallback, timeout: int) -> str:
        end_callback.set_timeout(timeout)
        self._end_callback = end_callback

        event_id = IdGenerator.generate_event_id()
        envelope = {
            "type": AgentEventType.AGENTSDK_FLOW_END_REQUEST.value,
            "agentId": self._agent_id,
            "agentMode": AgentPolicy.ACTIVE.value,
            "flowId": self._flow_id,
            "eventId": event_id,
        }
        text = json.dumps(envelope, ensure_ascii=False)
        self._send(text, event_id, lambda code, msg: end_callback.on_end_request_ack(self._flow_id, event_id, code, msg))
        return event_id

    def interrupt_request(
        self, interrupt_callback: FlowInterruptCallback, timeout: int
    ) -> str:
        interrupt_callback.set_timeout(timeout)
        self._interrupt_callback = interrupt_callback

        event_id = IdGenerator.generate_event_id()
        envelope = {
            "type": AgentEventType.AGENTSDK_FLOW_INTERRUPT_REQUEST.value,
            "agentId": self._agent_id,
            "agentMode": AgentPolicy.ACTIVE.value,
            "flowId": self._flow_id,
            "eventId": event_id,
        }
        text = json.dumps(envelope, ensure_ascii=False)
        self._send(text, event_id, lambda code, msg: interrupt_callback.on_interrupt_request_ack(self._flow_id, event_id, code, msg))
        return event_id

    # ----- internal -----

    def _send(self, text: str, event_id: str, ack_on_closed) -> None:
        client = self._linksky_client
        if client is None or client.is_closed():
            _log.error("flow %s send => linksky client is closed", self._flow_id)
            ack_on_closed(_CONN_CLOSED_CODE, _CONN_CLOSED_MSG)
        else:
            client.send_text(self._agent_id, self._flow_id, text, text)
        self._update_ts = int(time.time() * 1000)
