app/backend/app/flow/nodes/influx.py

"""InfluxDB nodes: write points from messages, or read a query into them."""

from __future__ import annotations

import logging
from collections.abc import Iterable
from typing import Any

from pydantic import BaseModel, ConfigDict, Field

from app.flow.messages import MessageSpec
from app.flow.nodes.base import Node

logger = logging.getLogger(__name__)


class InfluxDbNode(Node):
    """
    InfluxDB node for writing to and reading from InfluxDB.

    This node can perform both write and read operations independently:

    - **Write operation**: Triggered when upstream dependencies are satisfied
      (data flows in via ``requires``). Writes data points to InfluxDB based
      on the ``writes`` configuration in params.
    - **Read operation**: Performed when the node provides data to downstream
      nodes via ``provides``, based on the ``queries`` configuration in params.

    Both operations use a similar configuration pattern in params, making the
    API consistent and the input/output data simple (just values).

    :param requires: Messages to write to InfluxDB. The actual value from each
        message is written according to the corresponding config in ``writes``.
    :type requires: MessageSpec | list[MessageSpec] | None
    :param provides: Messages to read from InfluxDB. Each message gets its value
        from a query defined in ``queries``.
    :type provides: MessageSpec | list[MessageSpec] | None
    :param params: Parameters dict containing:
        - ``url`` (str): InfluxDB server URL (required)
        - ``token`` (str): Authentication token (required)
        - ``org`` (str): Organization name (required)
        - ``bucket`` (str): Bucket name (required)
        - ``write_precision`` (str): Write precision ("ns", "us", "ms", "s"), default "ms"
        - ``query_range`` (str): Default time range for queries, e.g., "-1h", "-24h"
        - ``writes`` (dict): Write configurations keyed by message name, each with:
            - ``measurement`` (str): Measurement name to write to
            - ``field`` (str): Field name to write (default: "value")
            - ``tags`` (dict): Static tags to add to each point
        - ``queries`` (dict): Query configurations keyed by message name, each with:
            - ``measurement`` (str): Measurement name to query
            - ``field`` (str): Field name to retrieve (default: "value")
            - ``tags`` (dict): Optional tag filters
            - ``range`` (str): Optional time range override
            - ``aggregation`` (str): Aggregation function ("mean", "last", "first", "max", "min")
    :type params: dict
    :param name: Optional name for the node.
    :type name: str | None

    :raises ValueError: If required params are missing or both requires and provides are empty.

    :example:
        Write-only node (writes temperature values):

        >>> writer = InfluxDbNode(
        ...     requires=[MessageSpec(name="temperature", dtype=DType.FLOAT)],
        ...     params={
        ...         "url": "http://localhost:8086",
        ...         "token": "my-token",
        ...         "org": "my-org",
        ...         "bucket": "sensors",
        ...         "writes": {
        ...             "temperature": {
        ...                 "measurement": "environment",
        ...                 "field": "temp_celsius",
        ...                 "tags": {"location": "room1", "sensor": "dht22"},
        ...             }
        ...         },
        ...     },
        ... )

        Read-only node (queries average temperature):

        >>> reader = InfluxDbNode(
        ...     provides=[MessageSpec(name="avg_temperature", dtype=DType.FLOAT)],
        ...     params={
        ...         "url": "http://localhost:8086",
        ...         "token": "my-token",
        ...         "org": "my-org",
        ...         "bucket": "sensors",
        ...         "queries": {
        ...             "avg_temperature": {
        ...                 "measurement": "environment",
        ...                 "field": "temp_celsius",
        ...                 "tags": {"location": "room1"},
        ...                 "range": "-1h",
        ...                 "aggregation": "mean",
        ...             }
        ...         },
        ...     },
        ... )

        Combined read/write node:

        >>> node = InfluxDbNode(
        ...     requires=[MessageSpec(name="raw_temp", dtype=DType.FLOAT)],
        ...     provides=[MessageSpec(name="avg_temp", dtype=DType.FLOAT)],
        ...     params={
        ...         "url": "http://localhost:8086",
        ...         "token": "my-token",
        ...         "org": "my-org",
        ...         "bucket": "sensors",
        ...         "writes": {
        ...             "raw_temp": {
        ...                 "measurement": "temperature",
        ...                 "field": "value",
        ...                 "tags": {"source": "sensor"},
        ...             }
        ...         },
        ...         "queries": {
        ...             "avg_temp": {
        ...                 "measurement": "temperature",
        ...                 "field": "value",
        ...                 "aggregation": "mean",
        ...                 "range": "-5m",
        ...             }
        ...         },
        ...     },
        ... )
    """

    __slots__ = (
        "url",
        "token",
        "org",
        "bucket",
        "write_precision",
        "query_range",
        "writes",
        "queries",
        "_write_client",
        "_query_client",
    )

    class Params(BaseModel):
        model_config = ConfigDict(extra="allow")

        url: str
        token: str = Field(json_schema_extra={"x-secret": True})
        org: str
        bucket: str
        write_precision: str = "ms"
        query_range: str = "-1h"
        # Per-port write and query configuration.
        writes: dict[str, dict[str, Any]] = {}
        queries: dict[str, dict[str, Any]] = {}

    def __init__(
        self,
        requires: MessageSpec | Iterable[MessageSpec] = (),
        provides: MessageSpec | Iterable[MessageSpec] = (),
        params: dict[str, Any] | None = None,
        name: str | None = None,
    ):
        cfg = self.Params.model_validate(params or {})

        requires = Node._normalize_ports(requires)
        provides = Node._normalize_ports(provides)
        if not requires and not provides:
            raise ValueError(
                "An InfluxDB node needs either inputs (to write) or outputs (to read)"
            )

        self.url = cfg.url
        self.token = cfg.token
        self.org = cfg.org
        self.bucket = cfg.bucket
        self.write_precision = cfg.write_precision
        self.query_range = cfg.query_range
        self.writes = cfg.writes
        self.queries = cfg.queries

        # Lazy-initialized clients
        self._write_client = None
        self._query_client = None

        # Set default name
        if name is None:
            name = f"influxdb_{self.bucket}"

        # Initialize parent
        # The handler function depends on what operations are configured
        super().__init__(
            f=self._handler,
            requires=requires,
            provides=provides,
            params=params,
            name=name,
        )

    def _handler(self, params: dict, **kwargs) -> dict | None:
        """
        Handle incoming data - write to InfluxDB and optionally query.

        This method is called when upstream dependencies (requires) are satisfied.
        It writes the incoming data to InfluxDB and can also perform reads.

        :param params: Node parameters.
        :type params: dict
        :param kwargs: Incoming data from upstream nodes.
        :type kwargs: Any
        :returns: Query results if provides is configured, None otherwise.
        :rtype: dict | None
        """
        # Write incoming data
        if kwargs:
            self._write_points(kwargs)

        # If we have provides, perform queries
        if self.provides:
            return self._query_data()

        return None

    def _write_points(self, data: dict) -> None:
        """
        Write data points to InfluxDB using configuration from ``writes``.

        The write configuration is looked up in ``self.writes`` by message name.
        Each config specifies the measurement, field, and tags. The actual value
        comes from the incoming data.

        :param data: Data to write, keyed by port name. Values can be:
            - Simple values (float, int, str, bool): Written using config from ``writes``
            - List of values: Each value written as a separate point
            - Dict with "value" key: Value extracted and written using config
            - Dict with "value" and "tags" keys: Value written with merged tags
        :type data: dict
        """
        from influxdb_client import InfluxDBClient, Point, WritePrecision
        from influxdb_client.client.write_api import SYNCHRONOUS

        try:
            with InfluxDBClient(url=self.url, token=self.token, org=self.org) as client:
                write_api = client.write_api(write_options=SYNCHRONOUS)

                precision_map = {
                    "ns": WritePrecision.NS,
                    "us": WritePrecision.US,
                    "ms": WritePrecision.MS,
                    "s": WritePrecision.S,
                }
                precision = precision_map.get(self.write_precision, WritePrecision.MS)

                for msg_name, msg_value in data.items():
                    # Get write configuration for this message
                    write_config = self.writes.get(msg_name, {})

                    # Get measurement, field, and base tags from config
                    measurement = write_config.get("measurement", msg_name)
                    field = write_config.get("field", "value")
                    base_tags = write_config.get("tags", {})

                    # Handle list of values (batch write)
                    values_to_write = (
                        msg_value if isinstance(msg_value, list) else [msg_value]
                    )

                    for item in values_to_write:
                        # Extract value and optional runtime tags
                        if isinstance(item, dict):
                            value = item.get("value", item)
                            runtime_tags = item.get("tags", {})
                            # If no "value" key, treat the whole dict as invalid
                            if "value" not in item and not isinstance(
                                value, (int, float, str, bool)
                            ):
                                logger.warning(
                                    "Skipping invalid item in node '%s': %s",
                                    self.name,
                                    item,
                                )
                                continue
                        else:
                            value = item
                            runtime_tags = {}

                        if value is None:
                            logger.info(
                                "Skipping None value for '%s' in node '%s'",
                                msg_name,
                                self.name,
                            )
                            continue

                        # Merge base tags with runtime tags (runtime takes precedence)
                        tags = {**base_tags, **runtime_tags}

                        # Build the point
                        point = Point(measurement)

                        for tag_key, tag_value in tags.items():
                            point = point.tag(tag_key, str(tag_value))

                        point = point.field(field, value)

                        # Write the point
                        write_api.write(
                            bucket=self.bucket,
                            org=self.org,
                            record=point,
                            write_precision=precision,
                        )

                        logger.info(
                            "Wrote to InfluxDB from node '%s': %s.%s=%s, tags=%s",
                            self.name,
                            measurement,
                            field,
                            value,
                            tags,
                        )

        except Exception as e:
            logger.error("InfluxDB write error in node '%s': %s", self.name, e)
            raise

    def _query_data(self) -> dict:
        """
        Query data from InfluxDB based on provides configuration.

        :returns: Dict of port name to queried value.
        :rtype: dict
        """
        from influxdb_client import InfluxDBClient

        results = {}

        try:
            with InfluxDBClient(url=self.url, token=self.token, org=self.org) as client:
                query_api = client.query_api()

                for spec in self.output_ports:
                    msg_name = spec.port
                    query_config = self.queries.get(msg_name, {})

                    measurement = query_config.get("measurement", msg_name)
                    field = query_config.get("field", "value")
                    tags = query_config.get("tags", {})
                    time_range = query_config.get("range", self.query_range)
                    aggregation = query_config.get("aggregation", "last")

                    # Build Flux query
                    flux_query = self._build_flux_query(
                        measurement=measurement,
                        field=field,
                        tags=tags,
                        time_range=time_range,
                        aggregation=aggregation,
                    )

                    logger.info(
                        "Executing InfluxDB query for '%s' in node '%s': %s",
                        msg_name,
                        self.name,
                        flux_query,
                    )

                    # Execute query
                    tables = query_api.query(flux_query, org=self.org)

                    # Extract result
                    value = self._extract_query_result(tables, spec)

                    if value is not None:
                        results[msg_name] = value
                        logger.info(
                            "Query result for '%s' in node '%s': %s",
                            msg_name,
                            self.name,
                            value,
                        )
                    else:
                        logger.info(
                            "No data found for '%s' in node '%s'",
                            msg_name,
                            self.name,
                        )

        except Exception as e:
            logger.error("InfluxDB query error in node '%s': %s", self.name, e)
            raise

        return results

    def _build_flux_query(
        self,
        measurement: str,
        field: str,
        tags: dict,
        time_range: str,
        aggregation: str,
    ) -> str:
        """
        Build a Flux query string.

        :param measurement: Measurement name.
        :type measurement: str
        :param field: Field name.
        :type field: str
        :param tags: Tag filters.
        :type tags: dict
        :param time_range: Time range (e.g., "-1h").
        :type time_range: str
        :param aggregation: Aggregation function.
        :type aggregation: str
        :returns: Flux query string.
        :rtype: str
        """
        # Base query
        query_parts = [
            f'from(bucket: "{self.bucket}")',
            f"  |> range(start: {time_range})",
            f'  |> filter(fn: (r) => r["_measurement"] == "{measurement}")',
            f'  |> filter(fn: (r) => r["_field"] == "{field}")',
        ]

        # Add tag filters
        for tag_key, tag_value in tags.items():
            query_parts.append(
                f'  |> filter(fn: (r) => r["{tag_key}"] == "{tag_value}")'
            )

        # Add aggregation
        aggregation_map = {
            "mean": "mean()",
            "last": "last()",
            "first": "first()",
            "max": "max()",
            "min": "min()",
            "sum": "sum()",
            "count": "count()",
        }

        if aggregation in aggregation_map:
            query_parts.append(f"  |> {aggregation_map[aggregation]}")
        else:
            # Default to last value
            query_parts.append("  |> last()")

        return "\n".join(query_parts)

    def _extract_query_result(self, tables, spec: MessageSpec) -> Any:
        """
        Extract a single value from query result tables.

        :param tables: InfluxDB query result tables.
        :param spec: The port the value is destined for.
        :type spec: MessageSpec
        :returns: Extracted and typed value, or None if no data.
        :rtype: Any
        """
        for table in tables:
            for record in table.records:
                value = record.get_value()
                if value is not None:
                    try:
                        return spec.coerce(value)
                    except (ValueError, TypeError):
                        return value

        return None

    def inject(self, outputs: dict | None = None) -> dict | None:
        """
        Inject queried data into the pipeline.

        For InfluxDbNode, inject performs a query operation and injects
        the results into the pipeline. This is useful for trigger-style
        usage where you want to periodically query InfluxDB.

        :param outputs: Optional pre-set outputs (usually None for queries).
        :type outputs: dict | None
        :returns: Query results injected into the pipeline.
        :rtype: dict | None
        """
        if self._pipeline is None:
            raise RuntimeError("Node must be bound to a pipeline to inject")

        # Without given values, injecting means running the configured queries.
        if not outputs:
            if not self.provides:
                return None
            outputs = self._query_data()

        return self._pipeline.trigger(self, self._to_messages(outputs))