AutoRpa/rpa_yisaiqi/callback_client.py

"""HTTP client for YiSaiQi callbacks into the state center."""
from __future__ import annotations

import json
from typing import Any, Callable, Dict

from urllib.error import HTTPError
from urllib.parse import quote
from urllib.request import Request, urlopen

JsonDict = Dict[str, Any]
Opener = Callable[[Request, float], Any]


class StateCenterError(RuntimeError):
    def __init__(self, status: int, detail: Any):
        super().__init__(f"state center returned {status}: {detail}")
        self.status = status
        self.detail = detail


class StateCenterClient:
    def __init__(
        self,
        base_url: str,
        token: str,
        *,
        timeout: float = 15,
        opener: Opener | None = None,
    ):
        self.base_url = base_url.rstrip("/")
        self.token = token
        self.timeout = timeout
        self._opener = opener or _urlopen_with_timeout

    def start_node(
        self,
        job_id: str,
        node_code: str,
        *,
        attempt_no: int = 1,
        operator: str = "agent",
    ) -> JsonDict:
        return self._post(
            f"/api/jobs/{_path_part(job_id)}/nodes/{_path_part(node_code)}/start",
            {
                "attemptNo": attempt_no,
                "executorPlatform": "iS-RPA",
                "operator": operator,
            },
        )

    def complete_node(
        self,
        job_id: str,
        node_code: str,
        payload: JsonDict,
    ) -> JsonDict:
        return self._post(f"/api/jobs/{_path_part(job_id)}/nodes/{_path_part(node_code)}/complete", payload)

    def fail_node(
        self,
        job_id: str,
        node_code: str,
        payload: JsonDict,
    ) -> JsonDict:
        return self._post(f"/api/jobs/{_path_part(job_id)}/nodes/{_path_part(node_code)}/fail", payload)

    def get_job(self, job_id: str) -> JsonDict:
        return self._request("GET", f"/api/jobs/{_path_part(job_id)}")

    def ingest_items(self, job_id: str, rows: list[JsonDict], replace: bool = True) -> JsonDict:
        return self._post(
            f"/api/jobs/{_path_part(job_id)}/items",
            {"replace": replace, "rows": rows},
        )

    def match_run(self, job_id: str) -> JsonDict:
        return self._post(f"/api/jobs/{_path_part(job_id)}/match/run", {})

    def load_fixture_goods(self) -> JsonDict:
        return self._post("/api/dev/match/load-fixture-goods", {})

    def retry_job(self, job_id: str) -> JsonDict:
        return self._post(f"/api/jobs/{_path_part(job_id)}/retry", {})

    def list_tasks(self, *, role_code: str | None = None, status: str | None = None) -> JsonDict:
        query = []
        if role_code:
            query.append(f"roleCode={quote(role_code, safe='')}")
        if status:
            query.append(f"status={quote(status, safe='')}")
        suffix = f"?{'&'.join(query)}" if query else ""
        return self._request("GET", f"/api/tasks/my{suffix}")

    def submit_task(
        self,
        task_id: str,
        action: str,
        payload: JsonDict,
        submit_token: str,
        operator: str = "agent",
    ) -> JsonDict:
        return self._post(
            f"/api/tasks/{_path_part(task_id)}/submit",
            {
                "action": action,
                "payload": payload,
                "submitToken": submit_token,
                "operator": operator,
            },
        )

    def check_contract_result(
        self,
        job_id: str,
        approved: bool,
        *,
        operator: str = "agent",
        detail: str | None = None,
    ) -> JsonDict:
        return self._post(
            f"/api/jobs/{_path_part(job_id)}/nodes/CONTRACT_APPROVE_WAIT/check-result",
            {"approved": approved, "operator": operator, "detail": detail},
        )

    def create_job(self, source_type: str, order_date: str, supplier_name: str, created_by: str = "agent") -> JsonDict:
        return self._post(
            "/api/jobs",
            {
                "source_type": source_type,
                "order_date": order_date,
                "supplier_name": supplier_name,
                "created_by": created_by,
            },
        )

    def _post(self, path: str, payload: JsonDict) -> JsonDict:
        return self._request("POST", path, payload)

    def _request(self, method: str, path: str, payload: JsonDict | None = None) -> JsonDict:
        data = json.dumps(payload, ensure_ascii=False).encode("utf-8") if payload is not None else None
        headers = {"X-Agent-Token": self.token}
        if data is not None:
            headers["Content-Type"] = "application/json"
        request = Request(
            f"{self.base_url}{path}",
            data=data,
            headers=headers,
            method=method,
        )
        try:
            with self._opener(request, self.timeout) as response:
                raw = response.read().decode("utf-8")
        except HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")
            raise StateCenterError(exc.code, detail) from exc
        return json.loads(raw) if raw else {}


def _path_part(value: str) -> str:
    return quote(value, safe="")


def _urlopen_with_timeout(request: Request, timeout: float) -> Any:
    return urlopen(request, timeout=timeout)