data-collector-mcp/data_collector_mcp/collector_api.py

from __future__ import annotations

from typing import Any

from .auth import find_project_config, resolve_project_token
from .http_client import request_json
from .protocols import MODBUS_SPEC, S7_SPEC
from .protocols.modbus import MODBUS_POINT_TYPE_ALIASES, MODBUS_REGISTER_TYPE_ALIASES
from .protocols.s7 import S7_POINT_TYPE_ALIASES, S7_REGISTER_TYPE_ALIASES


def _merge_defaults(defaults: dict[str, Any], payload: dict[str, Any]) -> dict[str, Any]:
    merged = dict(defaults)
    merged.update(payload)
    return merged


def _require_non_empty_text(payload: dict[str, Any], field_name: str) -> str:
    value = str(payload.get(field_name) or "").strip()
    if not value:
        raise ValueError(f"payload.{field_name} is required")
    return value


def _require_present(payload: dict[str, Any], field_name: str) -> Any:
    if field_name not in payload or payload.get(field_name) is None:
        raise ValueError(f"payload.{field_name} is required")
    return payload[field_name]


def _normalize_modbus_device_payload(payload: dict[str, Any]) -> dict[str, Any]:
    normalized = dict(payload)

    normalized["name"] = _require_non_empty_text(normalized, "name")
    normalized["ip"] = _require_non_empty_text(normalized, "ip")
    _require_present(normalized, "device_type")
    _require_present(normalized, "port")
    _require_present(normalized, "slave_id")
    _require_present(normalized, "word_order")
    _require_present(normalized, "byte_order")
    address_base = _require_present(normalized, "address_base")
    normalized["address_offset"] = address_base
    normalized.pop("address_base", None)

    return normalized


def _normalize_modbus_device_edit_payload(payload: dict[str, Any]) -> dict[str, Any]:
    normalized = dict(payload)

    normalized["ori_id"] = _normalize_positive_int(_require_present(normalized, "ori_id"), "payload.ori_id")
    normalized["name"] = _require_non_empty_text(normalized, "name")

    if normalized.get("type") is None:
        normalized["type"] = _require_present(normalized, "device_type")
    try:
        normalized["type"] = int(normalized["type"])
    except Exception as exc:
        raise ValueError("payload.type must be one of 1, 2, 3, 4, 5") from exc
    if normalized["type"] not in {1, 2, 3, 4, 5}:
        raise ValueError("payload.type must be one of 1, 2, 3, 4, 5")
    normalized.pop("device_type", None)

    if normalized["type"] == 2:
        normalized["serial_port"] = _require_non_empty_text(normalized, "serial_port")
    else:
        normalized["ip"] = _require_non_empty_text(normalized, "ip")
        _require_present(normalized, "port")

    _require_present(normalized, "slave_id")
    _require_present(normalized, "word_order")
    _require_present(normalized, "byte_order")

    if "address_base" in normalized:
        normalized["address_offset"] = normalized["address_base"]
        normalized.pop("address_base", None)
    if "group_id" in normalized and "device_group_id" not in normalized:
        normalized["device_group_id"] = normalized["group_id"]
        normalized.pop("group_id", None)

    return normalized


def _normalize_modbus_point_payload(payload: dict[str, Any]) -> dict[str, Any]:
    normalized = dict(payload)

    normalized["name"] = _require_non_empty_text(normalized, "name")
    _require_present(normalized, "address")

    raw_type = _require_non_empty_text(normalized, "type")
    normalized_type = MODBUS_POINT_TYPE_ALIASES.get(raw_type)
    if normalized_type is None:
        normalized_type = MODBUS_POINT_TYPE_ALIASES.get(raw_type.lower())
    if normalized_type is None:
        raise ValueError(
            "payload.type is invalid; use one of bool, int16, uint16, int32, "
            "uint32, int64, uint64, float32, float64, or a documented alias "
            "such as SHORT, WORD, LONG, DWORD, FLOAT, DOUBLE"
        )
    normalized["type"] = normalized_type

    if normalized.get("func_code") is None:
        register_type = _require_non_empty_text(normalized, "register_type")
        func_code = MODBUS_REGISTER_TYPE_ALIASES.get(register_type)
        if func_code is None:
            func_code = MODBUS_REGISTER_TYPE_ALIASES.get(register_type.lower())
        if func_code is None:
            raise ValueError(
                "payload.register_type is invalid; use coil, discrete_input, "
                "holding_register, input_register, or func_code 1/2/3/4"
            )
        normalized["func_code"] = func_code
    else:
        try:
            normalized["func_code"] = int(str(normalized["func_code"]).strip())
        except Exception as exc:
            raise ValueError("payload.func_code must be one of 1, 2, 3, 4") from exc
        if normalized["func_code"] not in {1, 2, 3, 4}:
            raise ValueError("payload.func_code must be one of 1, 2, 3, 4")

    if "register_type" in normalized:
        normalized.pop("register_type")

    return normalized


def _normalize_modbus_point_edit_payload(payload: dict[str, Any]) -> dict[str, Any]:
    _require_present(payload, "ori_id")
    normalized = _normalize_modbus_point_payload(payload)
    normalized["ori_id"] = _normalize_positive_int(_require_present(normalized, "ori_id"), "payload.ori_id")
    return normalized


def _normalize_s7_device_payload(payload: dict[str, Any]) -> dict[str, Any]:
    normalized = dict(payload)

    normalized["name"] = _require_non_empty_text(normalized, "name")
    normalized["ip"] = _require_non_empty_text(normalized, "ip")
    _require_present(normalized, "rock")
    _require_present(normalized, "slot")
    if "group_id" in normalized and "device_group_id" not in normalized:
        normalized["device_group_id"] = normalized["group_id"]
        normalized.pop("group_id", None)
    if "tsap_conn_type" in normalized:
        normalized["tsap_conn_type"] = _normalize_s7_tsap_conn_type(normalized["tsap_conn_type"])

    return normalized


def _normalize_s7_device_edit_payload(payload: dict[str, Any]) -> dict[str, Any]:
    normalized = _normalize_s7_device_payload(payload)
    if normalized.get("id") is None:
        normalized["id"] = _require_present(normalized, "ori_id")
    normalized["id"] = _normalize_positive_int(normalized["id"], "payload.id")
    normalized.pop("ori_id", None)
    return normalized


def _normalize_s7_point_payload(payload: dict[str, Any], *, require_device_id: bool = True) -> dict[str, Any]:
    normalized = dict(payload)

    normalized["name"] = _require_non_empty_text(normalized, "name")
    normalized["address"] = _require_non_empty_text(normalized, "address")
    if require_device_id:
        normalized["device_id"] = _normalize_positive_int(
            _require_present(normalized, "device_id"),
            "payload.device_id",
        )

    raw_type = normalized.get("data_type")
    if raw_type is None:
        raw_type = _require_non_empty_text(normalized, "type")
    else:
        raw_type = str(raw_type).strip()
        if not raw_type:
            raise ValueError("payload.data_type is required")
    normalized_type = S7_POINT_TYPE_ALIASES.get(raw_type)
    if normalized_type is None:
        normalized_type = S7_POINT_TYPE_ALIASES.get(raw_type.lower())
    if normalized_type is None:
        raise ValueError(
            "payload.data_type is invalid; use one of bool, uint8, int8, "
            "uint16, int16, uint32, int32, float32, float64, or a documented "
            "alias such as BOOL, BYTE, SINT, WORD, INT, DWORD, DINT, REAL, LREAL"
        )
    normalized["data_type"] = normalized_type
    normalized.pop("type", None)

    if normalized.get("register_type") is None:
        register_area = _require_non_empty_text(normalized, "register_area")
        register_type = S7_REGISTER_TYPE_ALIASES.get(register_area)
        if register_type is None:
            register_type = S7_REGISTER_TYPE_ALIASES.get(register_area.lower())
        if register_type is None:
            raise ValueError("payload.register_area is invalid; use I, Q, M, DB, V, AI, or register_type 1..6")
        normalized["register_type"] = register_type
    else:
        try:
            normalized["register_type"] = int(str(normalized["register_type"]).strip())
        except Exception as exc:
            raise ValueError("payload.register_type must be one of 1, 2, 3, 4, 5, 6") from exc
        if normalized["register_type"] not in {1, 2, 3, 4, 5, 6}:
            raise ValueError("payload.register_type must be one of 1, 2, 3, 4, 5, 6")
    normalized.pop("register_area", None)

    if "group_id" in normalized and "group_Id" not in normalized:
        normalized["group_Id"] = normalized["group_id"]
        normalized.pop("group_id", None)

    return normalized


def _normalize_s7_point_edit_payload(payload: dict[str, Any]) -> dict[str, Any]:
    normalized = _normalize_s7_point_payload(payload)
    if normalized.get("id") is None:
        normalized["id"] = _require_present(normalized, "ori_id")
    normalized["id"] = _normalize_positive_int(normalized["id"], "payload.id")
    normalized.pop("ori_id", None)
    return normalized


def _request_collector(
    project_key: str,
    method: str,
    path: str,
    *,
    json_payload: dict[str, Any] | None = None,
) -> dict[str, Any]:
    project = find_project_config(project_key)
    authorization = resolve_project_token(project)
    response_payload = request_json(
        method,
        f"{project['data_collector_base_url']}{path}",
        authorization,
        json_payload=json_payload,
    )
    if not isinstance(response_payload, dict):
        raise ValueError(f"collector API returned invalid payload for {path}: {response_payload}")
    return response_payload


def create_modbus_device(project_key: str, payload: dict[str, Any]) -> dict[str, Any]:
    return _request_collector(
        project_key,
        "POST",
        MODBUS_SPEC.create_device_path,
        json_payload=_merge_defaults(
            MODBUS_SPEC.device_defaults,
            _normalize_modbus_device_payload(payload),
        ),
    )


def create_modbus_point(project_key: str, payload: dict[str, Any]) -> dict[str, Any]:
    return _request_collector(
        project_key,
        "POST",
        MODBUS_SPEC.create_point_path,
        json_payload=_merge_defaults(
            MODBUS_SPEC.point_defaults,
            _normalize_modbus_point_payload(payload),
        ),
    )


def create_s7_device(project_key: str, payload: dict[str, Any]) -> dict[str, Any]:
    return _request_collector(
        project_key,
        "POST",
        S7_SPEC.create_device_path,
        json_payload=_merge_defaults(
            S7_SPEC.device_defaults,
            _normalize_s7_device_payload(payload),
        ),
    )


def create_s7_point(project_key: str, payload: dict[str, Any]) -> dict[str, Any]:
    return _request_collector(
        project_key,
        "POST",
        S7_SPEC.create_point_path,
        json_payload=_merge_defaults(
            S7_SPEC.point_defaults,
            _normalize_s7_point_payload(payload),
        ),
    )


def edit_modbus_device(project_key: str, payload: dict[str, Any]) -> dict[str, Any]:
    return _request_collector(
        project_key,
        "POST",
        "/api/collector/modbus/device/edit",
        json_payload=_merge_defaults(
            {
                "timeout": 3,
                "is_persistent": True,
                "device_group_id": 0,
                "alarm_interval": 90,
                "collect_interval": 5,
                "address_offset": 0,
                "retry_times": 0,
                "mode": 0,
            },
            _normalize_modbus_device_edit_payload(payload),
        ),
    )


def edit_modbus_point(project_key: str, payload: dict[str, Any]) -> dict[str, Any]:
    return _request_collector(
        project_key,
        "POST",
        "/api/collector/modbus/point/edit_collect_point",
        json_payload=_merge_defaults(
            MODBUS_SPEC.point_defaults,
            _normalize_modbus_point_edit_payload(payload),
        ),
    )


def edit_s7_device(project_key: str, payload: dict[str, Any]) -> dict[str, Any]:
    return _request_collector(
        project_key,
        "POST",
        "/api/collector/s7/device/update",
        json_payload=_merge_defaults(
            S7_SPEC.device_defaults,
            _normalize_s7_device_edit_payload(payload),
        ),
    )


def edit_s7_point(project_key: str, payload: dict[str, Any]) -> dict[str, Any]:
    return _request_collector(
        project_key,
        "POST",
        "/api/collector/s7/point/update",
        json_payload=_merge_defaults(
            S7_SPEC.point_defaults,
            _normalize_s7_point_edit_payload(payload),
        ),
    )


def list_devices(project_key: str, num_points: bool = False) -> dict[str, Any]:
    num_points_text = "true" if num_points else "false"
    return _request_collector(
        project_key,
        "GET",
        f"/api/collector/device?num_points={num_points_text}",
    )


def connect_device(project_key: str, device_id: int, device_type: str = "modbus") -> dict[str, Any]:
    return _set_device_connect_status(
        project_key,
        device_id=device_id,
        device_type=device_type,
        status=2,
    )


def disconnect_device(project_key: str, device_id: int, device_type: str = "modbus") -> dict[str, Any]:
    return _set_device_connect_status(
        project_key,
        device_id=device_id,
        device_type=device_type,
        status=1,
    )


def list_device_points(
    project_key: str,
    device_id: int,
    device_type: str = "modbus",
    group_id: int = 0,
) -> dict[str, Any]:
    return _request_collector(
        project_key,
        "POST",
        "/api/collector/common/device/get_collect_point",
        json_payload={
            "id": _normalize_positive_int(device_id, "device_id"),
            "type": _normalize_device_type(device_type),
            "group_id": _normalize_non_negative_int(group_id, "group_id"),
        },
    )


def _set_device_connect_status(
    project_key: str,
    *,
    device_id: int,
    device_type: str,
    status: int,
) -> dict[str, Any]:
    return _request_collector(
        project_key,
        "POST",
        "/api/collector/common/device/set_connect_status",
        json_payload={
            "id": _normalize_positive_int(device_id, "device_id"),
            "type": _normalize_device_type(device_type),
            "status": status,
        },
    )


def _normalize_device_type(device_type: str) -> str:
    normalized = str(device_type or "").strip()
    if not normalized:
        raise ValueError("device_type is required")
    return normalized


def _normalize_s7_tsap_conn_type(value: Any) -> str:
    normalized = str(value or "").strip().upper()
    if normalized not in {"PG", "OP", "BASIC"}:
        raise ValueError("payload.tsap_conn_type must be one of PG, OP, BASIC")
    return normalized


def _normalize_positive_int(value: int, field_name: str) -> int:
    try:
        normalized = int(value)
    except Exception as exc:
        raise ValueError(f"{field_name} must be a positive integer") from exc
    if normalized <= 0:
        raise ValueError(f"{field_name} must be a positive integer")
    return normalized


def _normalize_non_negative_int(value: int, field_name: str) -> int:
    try:
        normalized = int(value)
    except Exception as exc:
        raise ValueError(f"{field_name} must be a non-negative integer") from exc
    if normalized < 0:
        raise ValueError(f"{field_name} must be a non-negative integer")
    return normalized