from __future__ import annotations

import asyncio
import datetime
import re
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Literal, Union, cast, overload
from zoneinfo import ZoneInfo

from cognite.client._api_client import APIClient
from cognite.client.data_classes import SyntheticDatapoints, SyntheticDatapointsList, TimeSeries, TimeSeriesWrite
from cognite.client.data_classes.data_modeling.ids import NodeId
from cognite.client.utils._auxiliary import is_unlimited
from cognite.client.utils._concurrency import AsyncSDKTask, execute_async_tasks
from cognite.client.utils._identifier import Identifier, InstanceId
from cognite.client.utils._importing import local_import
from cognite.client.utils._time import convert_timezone_to_str, timestamp_to_ms
from cognite.client.utils.useful_types import SequenceNotStr

if TYPE_CHECKING:
    import sympy

    from cognite.client import AsyncCogniteClient
    from cognite.client.config import ClientConfig


def _supported_sympy_infix_ops(operation: type[sympy.Basic]) -> str | None:
    sympy = local_import("sympy")
    return {sympy.Add: "+", sympy.Mul: "*"}.get(operation)


def _supported_sympy_functions(operation: type[sympy.Basic]) -> str | None:
    sympy = local_import("sympy")
    return {
        sympy.cos: "cos",
        sympy.sin: "sin",
        sympy.sqrt: "sqrt",
        sympy.log: "ln",
        sympy.exp: "exp",
        sympy.Abs: "abs",
    }.get(operation)


class SyntheticDatapointsAPI(APIClient):
    _RESOURCE_PATH = "/timeseries/synthetic"

    def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: AsyncCogniteClient) -> None:
        super().__init__(config, api_version, cognite_client)
        self._DPS_LIMIT_SYNTH = 10_000

    def _get_semaphore(self, operation: Literal["read", "write", "delete"]) -> asyncio.BoundedSemaphore:
        from cognite.client import global_config

        return global_config.concurrency_settings.datapoints._semaphore_factory(
            operation, project=self._cognite_client.config.project
        )

    @overload
    async def query(
        self,
        expressions: SequenceNotStr[str] | SequenceNotStr[sympy.Basic],
        start: int | str | datetime.datetime,
        end: int | str | datetime.datetime,
        limit: int | None = None,
        variables: Mapping[str | sympy.Symbol, str | NodeId | TimeSeries | TimeSeriesWrite] | None = None,
        aggregate: str | None = None,
        granularity: str | None = None,
        target_unit: str | None = None,
        target_unit_system: str | None = None,
        timezone: str | datetime.timezone | ZoneInfo | None = None,
    ) -> SyntheticDatapointsList: ...

    @overload
    async def query(
        self,
        expressions: str | sympy.Basic,
        start: int | str | datetime.datetime,
        end: int | str | datetime.datetime,
        limit: int | None = None,
        variables: Mapping[str | sympy.Symbol, str | NodeId | TimeSeries | TimeSeriesWrite] | None = None,
        aggregate: str | None = None,
        granularity: str | None = None,
        target_unit: str | None = None,
        target_unit_system: str | None = None,
        timezone: str | datetime.timezone | ZoneInfo | None = None,
    ) -> SyntheticDatapoints: ...

    async def query(
        self,
        expressions: str | sympy.Basic | Sequence[str] | Sequence[sympy.Basic],
        start: int | str | datetime.datetime,
        end: int | str | datetime.datetime,
        limit: int | None = None,
        variables: Mapping[str | sympy.Symbol, str | NodeId | TimeSeries | TimeSeriesWrite] | None = None,
        aggregate: str | None = None,
        granularity: str | None = None,
        target_unit: str | None = None,
        target_unit_system: str | None = None,
        timezone: str | datetime.timezone | ZoneInfo | None = None,
    ) -> SyntheticDatapoints | SyntheticDatapointsList:
        """`Calculate the result of a function on time series. <https://api-docs.cognite.com/20230101/tag/Synthetic-Time-Series/operation/querySyntheticTimeseries>`_

        Info:
            You can read the guide to synthetic time series in our `documentation <https://docs.cognite.com/dev/concepts/resource_types/synthetic_timeseries>`_.

        Args:
            expressions (str | sympy.Basic | Sequence[str] | Sequence[sympy.Basic]): Functions to be calculated. Supports both strings and sympy expressions. Strings can have either the API `ts{}` syntax, or contain variable names to be replaced using the `variables` parameter.
            start (int | str | datetime.datetime): Inclusive start.
            end (int | str | datetime.datetime): Exclusive end.
            limit (int | None): Number of datapoints per expression to retrieve.
            variables (Mapping[str | sympy.Symbol, str | NodeId | TimeSeries | TimeSeriesWrite] | None): An optional map of symbol replacements.
            aggregate (str | None): use this aggregate when replacing entries from `variables`, does not affect time series given in the `ts{}` syntax.
            granularity (str | None): use this granularity with the aggregate.
            target_unit (str | None): use this target_unit when replacing entries from `variables`, does not affect time series given in the `ts{}` syntax.
            target_unit_system (str | None): Same as target_unit, but with unit system (e.g. SI). Only one of target_unit and target_unit_system can be specified.
            timezone (str | datetime.timezone | ZoneInfo | None): The timezone to use when aggregating datapoints. For aggregates of granularity 'hour' and longer,
                which time zone should we align to. Align to the start of the hour, start of the day or start of the month. For time zones of type Region/Location,
                the aggregate duration can vary, typically due to daylight saving time. For time zones of type UTC+/-HH:MM, use increments of 15 minutes. Default: "UTC" (None)

        Returns:
            SyntheticDatapoints | SyntheticDatapointsList: A SyntheticDatapointsList object containing the calculated data.

        Examples:

            Execute a synthetic time series query with an expression. Here we sum three time series plus a constant. The first is referenced by ID,
            the second by external ID, and the third by instance ID:

                >>> from cognite.client import CogniteClient, AsyncCogniteClient
                >>> client = CogniteClient()
                >>> # async_client = AsyncCogniteClient()  # another option
                >>> expression = '''
                ...     123
                ...     + ts{id:123}
                ...     + ts{externalId:'abc'}
                ...     + ts{space:'my-space',externalId:'my-ts-xid'}
                ... '''
                >>> dps = client.time_series.data.synthetic.query(
                ...     expressions=expression,
                ...     start="2w-ago",
                ...     end="now")

            You can also specify variables for an easier query syntax:

                >>> from cognite.client.data_classes.data_modeling.ids import NodeId
                >>> ts = client.time_series.retrieve(id=123)
                >>> variables = {
                ...     "A": ts,
                ...     "B": "my_ts_external_id",
                ...     "C": NodeId("my-space", "my-ts-xid"),
                ... }
                >>> dps = client.time_series.data.synthetic.query(
                ...     expressions="A+B+C", start="2w-ago", end="2w-ahead", variables=variables)

            Use sympy to build complex expressions:

                >>> from sympy import symbols, cos, sin
                >>> x, y = symbols("x y")
                >>> dps = client.time_series.data.synthetic.query(
                ...     [sin(x), y*cos(x)],
                ...     start="2w-ago",
                ...     end="now",
                ...     variables={x: "foo", y: "bar"},
                ...     aggregate="interpolation",
                ...     granularity="15m",
                ...     target_unit="temperature:deg_c",
                ...     timezone="Europe/Oslo",  # can also use this format: 'UTC+05:30'
                ... )
        """
        if is_unlimited(limit):
            limit = cast(int, float("inf"))

        if single_expr := not isinstance(expressions, SequenceNotStr):
            expressions = [cast(Union[str, "sympy.Basic"], expressions)]

        tasks = []
        for user_expr in cast(Sequence[Union[str, "sympy.Basic"]], expressions):
            expression, short_expression = self._build_expression(
                user_expr, variables, aggregate, granularity, target_unit, target_unit_system
            )
            query = {"expression": expression, "start": timestamp_to_ms(start), "end": timestamp_to_ms(end)}
            if timezone is not None:
                if isinstance(timezone, (ZoneInfo, datetime.timezone)):
                    timezone = convert_timezone_to_str(timezone)
                query["timeZone"] = timezone
            tasks.append(AsyncSDKTask(self._fetch_datapoints, query, limit, short_expression))

        datapoints_summary = await execute_async_tasks(tasks)
        datapoints_summary.raise_compound_exception_if_failed_tasks()
        return SyntheticDatapointsList(datapoints_summary.results) if not single_expr else datapoints_summary.results[0]

    async def _fetch_datapoints(self, query: dict[str, Any], limit: int, short_expression: str) -> SyntheticDatapoints:
        dps = None
        semaphore = self._get_semaphore("read")
        while True:
            query["limit"] = min(limit, self._DPS_LIMIT_SYNTH)
            resp = await self._post(
                url_path=self._RESOURCE_PATH + "/query", json={"items": [query]}, semaphore=semaphore
            )
            data = resp.json()["items"][0] | {"expression": short_expression, "timezone": query.get("timeZone")}
            new_dps = SyntheticDatapoints._load(data)
            if dps is None:
                dps = new_dps
            else:
                dps.timestamp.extend(new_dps.timestamp)
                dps.value.extend(new_dps.value)
                dps.error.extend(new_dps.error)
            limit -= (n_fetched := len(data["datapoints"]))
            if n_fetched < self._DPS_LIMIT_SYNTH or limit <= 0:
                break
            query["start"] = data["datapoints"][-1]["timestamp"] + 1
        return dps

    def _build_expression(
        self,
        expression: str | sympy.Basic,
        variables: Mapping[str | sympy.Symbol, str | NodeId | TimeSeries | TimeSeriesWrite] | None = None,
        aggregate: str | None = None,
        granularity: str | None = None,
        target_unit: str | None = None,
        target_unit_system: str | None = None,
    ) -> tuple[str, str]:
        if getattr(expression, "__sympy__", False) is True:
            if variables:
                expression_str = self._process_sympy_expression(cast("sympy.Basic", expression))
            else:
                raise ValueError(
                    "sympy expressions are only supported in combination with the `variables` parameter to map symbols to time series."
                )
        elif isinstance(expression, str):
            expression_str = expression
        else:
            raise TypeError(f"expression must be str or a sympy expression, not {type(expression)}")

        if aggregate and granularity:
            aggregate_str = f",aggregate:'{aggregate}',granularity:'{granularity}'"
        elif not aggregate and not granularity:
            aggregate_str = ""
        else:
            raise ValueError("Pass either both of 'aggregate' and 'granularity', or neither")

        target_unit_str = ""
        if target_unit:
            if target_unit_system:
                raise ValueError("Only one of 'target_unit' and 'target_unit_system' can be specified.")
            target_unit_str = f",targetUnit:'{target_unit}'"
        elif target_unit_system:
            target_unit_str = f",targetUnitSystem:'{target_unit_system}'"

        if not variables:
            return expression_str, expression_str

        to_substitute = {}
        for k, v in variables.items():
            if isinstance(v, TimeSeries | TimeSeriesWrite):
                try:
                    instance_id = getattr(v, "instance_id", None)
                    v = Identifier.load(external_id=v.external_id, instance_id=instance_id).as_primitive()
                except ValueError:
                    # ^error message wrongly says id is accepted, which it is not
                    raise ValueError(
                        f"TimeSeries passed in 'variables' is missing required field 'external_id' or 'instance_id' ({v})"
                    ) from None

            if isinstance(v, str):
                sub_string = f"externalId:'{v}'"
            elif isinstance(v, InstanceId):
                sub_string = f"space:'{v.space}',externalId:'{v.external_id}'"
            else:
                raise TypeError(
                    f"Unsupported variable type={type(v)} in 'variables', must be str, NodeId or TimeSeries object"
                )

            # We convert to str to ensure any sympy.Symbol is replaced with its name:
            to_substitute[re.escape(str(k))] = "ts{" + sub_string + aggregate_str + target_unit_str + "}"

        # Substitute all variables in one go to avoid substitution of prior substitutions:
        pattern = re.compile(r"\b" + r"\b|\b".join(to_substitute) + r"\b")  # note: \b marks a word boundary
        expression_with_ts = pattern.sub(lambda match: to_substitute[match[0]], expression_str)
        return expression_with_ts, expression_str

    def _process_sympy_expression(self, expression: sympy.Basic) -> str:
        sympy = local_import("sympy")

        if isinstance(expression, sympy.AtomicExpr):
            if isinstance(expression, sympy.NumberSymbol):
                return str(expression.evalf(15)).rstrip("0")
            else:
                return str(expression).rstrip("0")

        expr_cls = type(expression)
        if infix_op := _supported_sympy_infix_ops(expr_cls):
            return "(" + infix_op.join(self._process_sympy_expression(s) for s in expression.args) + ")"

        if isinstance(expression, sympy.Pow):
            if expression.args[1] == -1:
                return f"(1/{self._process_sympy_expression(expression.args[0])})"
            return f"pow({','.join(map(self._process_sympy_expression, expression.args))})"

        if fn_op := _supported_sympy_functions(expr_cls):
            return f"{fn_op}({','.join(map(self._process_sympy_expression, expression.args))})"
        raise TypeError(f"Unsupported sympy class {expr_cls} encountered in expression")
