"""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 fluksio.flow.messages import MessageSpec
from fluksio.flow.nodes.base import Node, NodeResult

logger = logging.getLogger(__name__)


def _influxdb() -> Any:
    """The client library, which is a `fluksio[server]` extra.

    Imported per use rather than at module level, because the node type is
    registered at boot and an instance with no InfluxDB behind it should
    not have to carry the library to start.
    """
    try:
        import influxdb_client
    except ImportError:
        raise RuntimeError(
            "the influxdb node needs the server extra: pip install 'fluksio[server]'"
        ) from None
    return influxdb_client


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).

    - **Query passthrough**: an incoming message holding a ``flux`` key is run
      as written, and every row comes back on the first output port as
      ``{"rows": [{ts, value, field, measurement, tags}], **echo}``.

    The passthrough is what keeps a database node a database node. It holds the
    credentials and the connection and nothing else: building a query and
    shaping its rows are ordinary Python nodes on either side, so a dashboard
    widget never learns which database answered it. A chart drawing an
    InfluxDB series is::

        [chart] -record-> [py: build] -record{flux, ...}-> [influx]
        [chart] <-series- [py: parse] <-json{rows, ...}---'

    Nothing rate-limits the requests here; an ``interval`` on the request port
    is what holds back a caller asking too often.

    :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"
        - ``timeout`` (float): Deadline for a request in seconds (default: 10.0)
        - ``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",
        ...             }
        ...         },
        ...     },
        ... )
    """

    # Writing the same point twice doubles it in the series.
    idempotent = False

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

    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"
        # Bounds every request to the server. Without one the client falls
        # back to its own default, which no flow can see or change. Held in
        # seconds like every other node; the client counts in milliseconds.
        timeout: float = Field(
            default=10.0,
            gt=0,
            description="Give up on a query or write after this many seconds.",
        )
        # Per-port write and query configuration.
        writes: dict[str, dict[str, Any]] = {}
        queries: dict[str, dict[str, Any]] = {}

    @classmethod
    def target_key(cls, params: dict[str, Any]) -> str | None:
        """The bucket, which is the thing several flows share."""
        url, bucket = params.get("url"), params.get("bucket")
        return f"{url}/{bucket}" if url and bucket else None

    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.timeout = cfg.timeout
        self.writes = cfg.writes
        self.queries = cfg.queries

        # 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[str, Any], **kwargs: Any) -> dict[str, Any] | None:
        """
        Handle incoming data - write to InfluxDB, run a query, or both.

        This method is called when upstream dependencies (requires) are satisfied.

        An incoming dict carrying a ``flux`` key is a query request rather than
        something to store: it is executed as written and answered on the first
        output port. Everything else on the inputs is written as points.

        :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
        """
        # ponytail: one request/answer pair per node — the first input holding
        # a "flux" key is the request and the answer goes out on the first
        # output port. A second query stream means a second node.
        request = next(
            (v for v in kwargs.values() if isinstance(v, dict) and "flux" in v), None
        )
        writes = {k: v for k, v in kwargs.items() if v is not request}

        if writes:
            self._write_points(writes)

        if request is not None:
            if not self.output_ports:
                raise ValueError(
                    f"Node '{self.name}' was sent a query but has no output to "
                    "answer on"
                )
            return {self.output_ports[0].port: self._run_flux(request)}

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

        return None

    def _run_flux(self, request: dict[str, Any]) -> dict[str, Any]:
        """
        Execute a Flux query exactly as it was handed over, and answer its rows.

        The node stays out of the query's business: building it is the job of
        whatever produced the request, and shaping the rows into something a
        widget draws is the job of whatever reads the answer. Every field of
        the request except ``flux`` is echoed back untouched, which is how a
        caller tells its own answer from someone else's.

        :param request: The query and whatever the caller wants echoed.
        :type request: dict
        :returns: ``{"rows": [...], **echo}``.
        :rtype: dict
        """
        InfluxDBClient = _influxdb().InfluxDBClient

        flux = str(request["flux"])
        echo = {key: value for key, value in request.items() if key != "flux"}

        logger.info("Running Flux for node '%s': %s", self.name, flux)
        with InfluxDBClient(
            url=self.url,
            token=self.token,
            org=self.org,
            timeout=int(self.timeout * 1000),
        ) as client:
            tables = client.query_api().query(flux, org=self.org)

        rows = [
            {
                "ts": record.get_time().timestamp() if record.get_time() else None,
                "value": record.get_value(),
                # Read off `values` rather than the getters: a query that keeps
                # only what it needs drops these columns, and the getters raise.
                "field": record.values.get("_field"),
                "measurement": record.values.get("_measurement"),
                "tags": {
                    key: value
                    for key, value in record.values.items()
                    if not key.startswith("_") and key not in ("result", "table")
                },
            }
            for table in tables
            for record in table.records
        ]
        return {"rows": rows, **echo}

    def _write_points(self, data: dict[str, Any]) -> 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
        """
        influxdb_client = _influxdb()
        InfluxDBClient = influxdb_client.InfluxDBClient
        Point, WritePrecision = influxdb_client.Point, influxdb_client.WritePrecision
        SYNCHRONOUS = influxdb_client.client.write_api.SYNCHRONOUS

        try:
            with InfluxDBClient(
                url=self.url,
                token=self.token,
                org=self.org,
                timeout=int(self.timeout * 1000),
            ) 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[str, Any]:
        """
        Query data from InfluxDB based on provides configuration.

        :returns: Dict of port name to queried value.
        :rtype: dict
        """
        InfluxDBClient = _influxdb().InfluxDBClient

        results = {}

        try:
            with InfluxDBClient(
                url=self.url,
                token=self.token,
                org=self.org,
                timeout=int(self.timeout * 1000),
            ) 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[str, Any],
        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: Any, 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[str, Any] | None = None, durable: bool | None = None
    ) -> NodeResult:
        """
        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), durable=durable)