"""
CDF SDK function wrappers for cdf_sdk: SPARQL namespace.

Uses instance_id (space + external_id) for all time series operations,
which is the recommended approach for DMS-integrated time series.

Functions available:
- cdf_sdk:datapoints_aggregate - Aggregate datapoints over time range
- cdf_sdk:datapoints_count - Count datapoints in time range
- cdf_sdk:datapoints_latest - Get latest datapoint value
- cdf_sdk:timeseries_exists - Check if time series exists

Reference:
- Cognite SDK Datapoints API: https://cognite-sdk-python.readthedocs-hosted.com/en/latest/time_series.html
"""

from __future__ import annotations

import logging
import time
from collections.abc import Callable
from typing import TYPE_CHECKING

from rdflib import Literal

from ._helpers import literal_to_python, parse_instance_id_from_uri, safe_sparql_wrapper

if TYPE_CHECKING:
    from cognite.client import CogniteClient

logger = logging.getLogger(__name__)

# Module-level prefetch caches (populated by prefetch_timeseries_cache before pyshacl runs)
_PREFETCH_EXISTS: dict[str, bool] = {}  # key: "space/external_id"
_PREFETCH_COUNTS: dict[str, int] = {}  # key: "space/external_id|start|end|granularity"
_PREFETCH_ERRORS: set[str] = set()  # keys where batch API permanently failed (all retries exhausted)
_COUNT_ERROR_SENTINEL = -1  # stored in _PREFETCH_COUNTS for errored entries (distinct from 0 = no datapoints)


def _retry_on_5xx(fn: Callable, max_retries: int = 3, base_delay: float = 1.0) -> object:
    """Call fn(), retrying up to max_retries times on 5xx CogniteAPIError with exponential backoff."""
    from cognite.client.exceptions import CogniteAPIError

    last_exc: Exception | None = None
    for attempt in range(max_retries + 1):
        try:
            return fn()
        except CogniteAPIError as e:
            if e.code is not None and 500 <= e.code < 600:
                last_exc = e
                if attempt < max_retries:
                    delay = base_delay * (2**attempt)
                    logger.warning(
                        f"[prefetch] 5xx error (HTTP {e.code}), retrying in {delay:.1f}s"
                        f" (attempt {attempt + 1}/{max_retries})..."
                    )
                    time.sleep(delay)
            else:
                raise
    raise last_exc  # type: ignore[misc]


def extract_count_prefetch_spec(
    shacl_graph: object,
    data_graph: object,
) -> tuple[list[str], list[tuple[str, str]]]:
    """
    Inspect the SHACL graph to find which time series need count prefetching and for what windows.

    Scans sh:select text in SHACL SPARQL constraints for calls to datapoints_count($this, start, end),
    resolves the sh:targetClass for each shape, finds instances of those classes in data_graph,
    and returns the unique (start, end) windows.

    Args:
        shacl_graph: rdflib Graph containing parsed SHACL rules
        data_graph: rdflib Graph containing instance data

    Returns:
        (uris, windows) where:
            uris    — deduplicated list of subject URIs in data_graph that are instances
                      of a target class for a count-using shape
            windows — deduplicated list of (start, end) string pairs extracted from the rules
    """
    import re

    from rdflib import RDF, SH, URIRef
    from rdflib import Graph as RDFGraph

    sg: RDFGraph = shacl_graph  # type: ignore[assignment]
    dg: RDFGraph = data_graph  # type: ignore[assignment]

    # Regex: datapoints_count($this, "start", "end")  — arguments are quoted strings
    count_pattern = re.compile(r'datapoints_count\(\s*\$this\s*,\s*"([^"]+)"\s*,\s*"([^"]+)"\s*\)')

    SH_SELECT = SH.select
    SH_TARGET_CLASS = SH.targetClass
    SH_SPARQL = SH.sparql

    windows_found: list[tuple[str, str]] = []
    target_classes: set[str] = set()

    # Walk: NodeShape → sh:sparql → sh:select
    for shape in sg.subjects(RDF.type, SH.NodeShape):
        for constraint in sg.objects(shape, SH_SPARQL):
            select_text = sg.value(constraint, SH_SELECT)
            if select_text is None:
                continue
            for m in count_pattern.finditer(str(select_text)):
                start, end = m.group(1), m.group(2)
                if (start, end) not in windows_found:
                    windows_found.append((start, end))
                # Collect target classes for this shape
                for tc in sg.objects(shape, SH_TARGET_CLASS):
                    target_classes.add(str(tc))

    if not windows_found or not target_classes:
        return [], windows_found

    # Find data_graph subjects that are instances of any of those target classes
    from rdflib import BNode

    uris: list[str] = []
    seen: set[str] = set()
    for tc_uri in target_classes:
        for subj in dg.subjects(RDF.type, URIRef(tc_uri)):
            if not isinstance(subj, BNode):
                s = str(subj)
                if s not in seen:
                    seen.add(s)
                    uris.append(s)

    return uris, windows_found


def get_prefetch_errors() -> set[str]:
    """Return copy of keys (space/external_id) where batch API calls permanently failed."""
    return set(_PREFETCH_ERRORS)


def clear_prefetch_cache() -> None:
    """Clear all prefetch caches (call before each validation run)."""
    _PREFETCH_EXISTS.clear()
    _PREFETCH_COUNTS.clear()
    _PREFETCH_ERRORS.clear()


def prefetch_timeseries_cache(
    client: CogniteClient,
    uris: list[str],
    count_uris: list[str] | None = None,
    count_windows: list[tuple[str, str]] | None = None,
    granularity: str = "1h",
    batch_size: int = 100,
    verbose: bool = False,
) -> None:
    """
    Batch pre-fetch time series existence and datapoint counts before pyshacl runs.

    Phase 1 (existence): runs for all URIs in `uris`.
    Phase 2 (counts): runs only for URIs in `count_uris` that were found to exist,
    and only for the (start, end) windows extracted from SHACL rules.
    If count_windows or count_uris is empty/None, Phase 2 is skipped entirely.

    Args:
        client: CogniteClient instance
        uris: All subject URIs from the data graph — used for Phase 1 (existence check).
        count_uris: Subset of URIs targeted by datapoints_count SHACL shapes. Only these
                    are included in the count pre-fetch (Phase 2). Use extract_count_prefetch_spec.
        count_windows: List of (start, end) pairs from SHACL datapoints_count rules.
                       If None or empty, no count prefetching is done.
        granularity: Granularity for count aggregation (default "1h")
        batch_size: Time series per API batch (default 100)
        verbose: Print progress
    """
    from cognite.client.data_classes.data_modeling import NodeId

    # Parse URIs to instance_ids, deduplicate.
    # Skip system spaces (e.g. cdf_cdm, cdf_system): CDF DMS forbids mixing
    # system spaces with user spaces in a single query, causing 400 errors.
    # CDF guarantees user-defined spaces cannot start with "cdf_".
    seen_keys: set[str] = set()
    instance_ids: list[NodeId] = []

    for uri in uris:
        try:
            iid = parse_instance_id_from_uri(uri)
            if iid.space.startswith("cdf_"):
                continue
            key = f"{iid.space}/{iid.external_id}"
            if key not in seen_keys:
                seen_keys.add(key)
                instance_ids.append(iid)
        except Exception:
            pass

    if not instance_ids:
        return

    if verbose:
        print(f"  [prefetch] Batch-checking existence of {len(instance_ids)} time series...")

    # Phase 1: Batch-check existence via retrieve_multiple
    existing_ts: dict[str, object] = {}  # cache_key → TimeSeries object
    failed_exist_keys: set[str] = set()  # keys from permanently failed batches

    for i in range(0, len(instance_ids), batch_size):
        batch = instance_ids[i : i + batch_size]
        try:
            ts_list = _retry_on_5xx(
                lambda b=batch: client.time_series.retrieve_multiple(
                    instance_ids=b,
                    ignore_unknown_ids=True,
                )
            )
            for ts in ts_list:
                if ts.instance_id:
                    key = f"{ts.instance_id.space}/{ts.instance_id.external_id}"
                    existing_ts[key] = ts
                elif ts.external_id:
                    # Fallback: match by external_id
                    for iid in batch:
                        if ts.external_id == iid.external_id:
                            key = f"{iid.space}/{iid.external_id}"
                            existing_ts[key] = ts
                            break
        except Exception as e:
            logger.warning(f"[prefetch] exists batch {i // batch_size + 1} failed: {e}")
            for iid in batch:
                failed_exist_keys.add(f"{iid.space}/{iid.external_id}")

    _PREFETCH_ERRORS.update(failed_exist_keys)

    # Populate exists cache — skip errored keys (their state is unknown, not False)
    for iid in instance_ids:
        key = f"{iid.space}/{iid.external_id}"
        if key not in _PREFETCH_ERRORS:
            _PREFETCH_EXISTS[key] = key in existing_ts

    if verbose:
        print(f"  [prefetch] {len(existing_ts)}/{len(instance_ids)} time series exist in CDF")

    if not existing_ts:
        return

    # Phase 2: Batch-fetch datapoint counts — only for windows extracted from SHACL rules,
    # and only for URIs targeted by count shapes (count_uris).
    # Using instance_id (not external_id) stays on the DMS query path (same restriction
    # as Phase 1; cdf_* spaces already filtered above).
    if not count_windows or not count_uris:
        if verbose:
            print("  [prefetch] No datapoints_count rules/targets found in SHACL — skipping count pre-fetch")
        return

    # Resolve count_uris → set of "space/external_id" keys (decoded, same format as existing_ts keys)
    count_keys: set[str] = set()
    for uri in count_uris:
        try:
            iid = parse_instance_id_from_uri(uri)
            if not iid.space.startswith("cdf_"):
                count_keys.add(f"{iid.space}/{iid.external_id}")
        except Exception:
            pass

    # Build NodeId list: only existing time series that are also in count_keys
    existing_node_ids: list[tuple[str, NodeId]] = []  # (ts_key "space/ext_id", NodeId)
    for key, ts in existing_ts.items():
        if key not in count_keys:
            continue
        iid = ts.instance_id
        if iid is None and ts.external_id:
            space, ext_id = key.split("/", 1)
            from cognite.client.data_classes.data_modeling import NodeId as _NodeId

            iid = _NodeId(space=space, external_id=ext_id)
        if iid is not None:
            existing_node_ids.append((key, iid))

    if not existing_node_ids:
        return

    total_counts_fetched = 0

    for start, end in count_windows:
        count_key_suffix = f"|{start}|{end}|{granularity}"
        window_items: list[tuple[str, NodeId]] = [
            (f"{ts_key}{count_key_suffix}", iid) for ts_key, iid in existing_node_ids
        ]

        if verbose:
            print(
                f"  [prefetch] Batch-fetching datapoint counts for {len(window_items)} time series"
                f" (start={start!r}, end={end!r}, granularity={granularity!r})..."
            )

        for i in range(0, len(window_items), batch_size):
            batch_items = window_items[i : i + batch_size]
            batch_node_ids = [iid for _, iid in batch_items]

            try:
                results = _retry_on_5xx(
                    lambda nids=batch_node_ids, s=start, e=end: client.time_series.data.retrieve(
                        instance_id=nids,
                        aggregates=["count"],
                        granularity=granularity,
                        start=s,
                        end=e,
                        ignore_unknown_ids=True,
                    )
                )

                # Build "space/external_id" → total count mapping from instance_id on results
                key_to_count: dict[str, int] = {}
                if results is not None:
                    dp_list = results if hasattr(results, "__iter__") else [results]
                    for dp_item in dp_list:
                        if dp_item.instance_id:
                            k = f"{dp_item.instance_id.space}/{dp_item.instance_id.external_id}"
                        elif dp_item.external_id:
                            k = dp_item.external_id  # fallback if SDK doesn't echo instance_id
                        else:
                            continue
                        total = sum(dp.count for dp in dp_item if dp.count is not None)
                        key_to_count[k] = total

                for cache_key, iid in batch_items:
                    ts_key = f"{iid.space}/{iid.external_id}"
                    _PREFETCH_COUNTS[cache_key] = key_to_count.get(ts_key, 0)
                    total_counts_fetched += 1

            except Exception as e:
                logger.warning(f"[prefetch] counts batch {i // batch_size + 1} failed: {e}")
                for cache_key, iid in batch_items:
                    _PREFETCH_COUNTS[cache_key] = _COUNT_ERROR_SENTINEL
                    _PREFETCH_ERRORS.add(f"{iid.space}/{iid.external_id}")

    if verbose:
        n_with_data = sum(1 for v in _PREFETCH_COUNTS.values() if v > 0)
        print(f"  [prefetch] {n_with_data}/{total_counts_fetched} time series have data in fetched windows")


def _parse_time_param(value: str | int) -> str | int:
    """
    Parse time parameter, converting millisecond timestamp strings to integers.

    The CDF SDK accepts:
    - Relative strings: "30d-ago", "7d-ago", "now"
    - Integer milliseconds: 1736848800000

    But NOT string milliseconds: "1736848800000"

    This function converts numeric strings to integers for SDK compatibility.
    """
    if isinstance(value, str):
        # Check if it's a numeric string (milliseconds timestamp)
        if value.isdigit() or (value.startswith("-") and value[1:].isdigit()):
            return int(value)
    return value


def create_sdk_wrappers(client: CogniteClient, verbose: bool = False) -> dict[str, Callable]:
    """
    Create CDF SDK wrapper functions for SPARQL registration.

    All functions use instance_id (space + external_id) for time series
    identification, leveraging the existing session client.

    Args:
        client: CogniteClient for CDF operations (from session state)

    Returns:
        Dict mapping function names to wrapper functions.
        Keys are the local part of the SPARQL function URI.

    Example:
        >>> wrappers = create_sdk_wrappers(client)
        >>> wrappers["datapoints_aggregate"]("http://...uri...", "count", "1h")
        Literal(42)
    """
    wrappers: dict[str, Callable] = {}

    # 1. Datapoints Aggregate
    @safe_sparql_wrapper(default_value=Literal(0))
    def datapoints_aggregate(
        timeseries_uri: str,
        aggregate: str = "count",
        granularity: str = "1h",
        start: str = "30d-ago",
        end: str = "now",
    ) -> Literal:
        """
        Aggregate datapoints for a time series using SDK aggregation.

        Supported aggregates: count, average, min, max, sum, interpolation,
        step_interpolation, continuous_variance, discrete_variance, total_variation.

        Args:
            timeseries_uri: URI containing space and external_id
            aggregate: Aggregation type (e.g., count, average, min, max, sum)
            granularity: Time granularity (e.g., 1h, 1d, 1w)
            start: Start time (e.g., "30d-ago", "2024-01-01")
            end: End time (e.g., "now", "2024-12-31")

        Returns:
            Sum of aggregated values across all granularity buckets as Literal
        """
        # Convert rdflib types to Python
        timeseries_uri = literal_to_python(timeseries_uri)
        aggregate = literal_to_python(aggregate)
        granularity = literal_to_python(granularity)
        start = _parse_time_param(literal_to_python(start))
        end = _parse_time_param(literal_to_python(end))

        instance_id = parse_instance_id_from_uri(timeseries_uri)

        # Use SDK aggregation - returns Datapoints object with aggregate values
        result = client.time_series.data.retrieve(
            instance_id=instance_id,
            aggregates=[aggregate],
            granularity=granularity,
            start=start,
            end=end,
        )

        if result is None or len(result) == 0:
            return Literal(0)

        # Sum up all aggregate values across the granularity buckets
        total = 0.0
        for datapoint in result:
            value = getattr(datapoint, aggregate, None)
            if value is not None:
                total += float(value)

        return Literal(total)

    wrappers["datapoints_aggregate"] = datapoints_aggregate

    # 2. Datapoints Count (convenience wrapper using SDK aggregation)
    @safe_sparql_wrapper(default_value=Literal(0))
    def datapoints_count(
        timeseries_uri: str,
        start: str = "30d-ago",
        end: str = "now",
    ) -> Literal:
        """
        Count datapoints in a time range using SDK aggregation.

        Args:
            timeseries_uri: URI containing space and external_id
            start: Start time
            end: End time

        Returns:
            Total number of datapoints as Literal
        """
        timeseries_uri = literal_to_python(timeseries_uri)
        start = _parse_time_param(literal_to_python(start))
        end = _parse_time_param(literal_to_python(end))

        instance_id = parse_instance_id_from_uri(timeseries_uri)

        # Use SDK aggregation with count - use appropriate granularity based on time range
        # For short ranges (< 1 day), use hourly granularity
        granularity = "1h"
        if isinstance(start, int) and isinstance(end, int):
            range_ms = end - start
            if range_ms > 86400000:  # More than 1 day
                granularity = "1d"

        # Check prefetch cache (populated by prefetch_timeseries_cache before pyshacl runs)
        cache_key = f"{instance_id.space}/{instance_id.external_id}|{start}|{end}|{granularity}"
        total_count = _PREFETCH_COUNTS.get(cache_key, 0)
        if total_count == _COUNT_ERROR_SENTINEL:
            if verbose:
                print(f"    [cdf_sdk] datapoints_count({instance_id}, start={start!r}, end={end!r}) → -1 (fetch error)")
            return Literal(_COUNT_ERROR_SENTINEL)
        if verbose:
            print(
                f"    [cdf_sdk] datapoints_count({instance_id}, start={start!r}, end={end!r}) → {total_count} (cached)"
            )
        return Literal(total_count)

    wrappers["datapoints_count"] = datapoints_count

    # 3. Datapoints Latest
    @safe_sparql_wrapper(default_value=Literal(float("nan")))
    def datapoints_latest(timeseries_uri: str) -> Literal:
        """
        Get the latest datapoint value for a time series.

        Args:
            timeseries_uri: URI containing space and external_id

        Returns:
            Latest value as Literal, or NaN if no data
        """
        timeseries_uri = literal_to_python(timeseries_uri)
        instance_id = parse_instance_id_from_uri(timeseries_uri)

        result = client.time_series.data.retrieve_latest(
            instance_id=instance_id,
        )

        if result and len(result) > 0 and result[0].datapoints:
            latest = result[0].datapoints[-1]
            return Literal(latest.value)

        return Literal(float("nan"))

    wrappers["datapoints_latest"] = datapoints_latest

    # 4. Time Series Exists
    @safe_sparql_wrapper(default_value=Literal(False))
    def timeseries_exists(timeseries_uri: str) -> Literal:
        """
        Check if a time series exists using instance_id.

        Args:
            timeseries_uri: URI containing space and external_id

        Returns:
            True if exists, False otherwise
        """
        timeseries_uri = literal_to_python(timeseries_uri)
        instance_id = parse_instance_id_from_uri(timeseries_uri)

        # Check prefetch cache (populated by prefetch_timeseries_cache before pyshacl runs).
        # Check _PREFETCH_EXISTS first — Phase 1 (existence batch) is authoritative.
        # _PREFETCH_ERRORS may also contain this key if Phase 2 (counts batch) failed,
        # but that does not affect whether the time series exists.
        cache_key = f"{instance_id.space}/{instance_id.external_id}"
        if cache_key in _PREFETCH_EXISTS:
            exists = _PREFETCH_EXISTS[cache_key]
            if verbose:
                print(f"    [cdf_sdk] timeseries_exists({instance_id}) → {exists} (cached)")
            return Literal(exists)
        if cache_key in _PREFETCH_ERRORS:
            # Phase 1 itself failed for this key — assume exists to avoid false violations
            if verbose:
                print(f"    [cdf_sdk] timeseries_exists({instance_id}) → True (fetch error, assuming exists)")
            return Literal(True)
        if verbose:
            print(f"    [cdf_sdk] timeseries_exists({instance_id}) → False (not in prefetch cache)")
        return Literal(False)

    wrappers["timeseries_exists"] = timeseries_exists

    # 5. Datapoints Average (using SDK aggregation)
    @safe_sparql_wrapper(default_value=Literal(float("nan")))
    def datapoints_average(
        timeseries_uri: str,
        start: str = "30d-ago",
        end: str = "now",
    ) -> Literal:
        """
        Calculate weighted average value over time range using SDK aggregation.

        Args:
            timeseries_uri: URI containing space and external_id
            start: Start time
            end: End time

        Returns:
            Weighted average value as Literal
        """
        timeseries_uri = literal_to_python(timeseries_uri)
        start = _parse_time_param(literal_to_python(start))
        end = _parse_time_param(literal_to_python(end))

        instance_id = parse_instance_id_from_uri(timeseries_uri)

        # Use SDK aggregation with both average and count to compute weighted average
        # Use appropriate granularity based on time range
        granularity = "1h"
        if isinstance(start, int) and isinstance(end, int):
            range_ms = end - start
            if range_ms > 86400000:  # More than 1 day
                granularity = "1d"

        result = client.time_series.data.retrieve(
            instance_id=instance_id,
            aggregates=["average", "count"],
            granularity=granularity,
            start=start,
            end=end,
        )

        if result is None or len(result) == 0:
            return Literal(float("nan"))

        # Compute weighted average across buckets
        total_weighted = 0.0
        total_count = 0
        for datapoint in result:
            if datapoint.average is not None and datapoint.count is not None:
                total_weighted += datapoint.average * datapoint.count
                total_count += datapoint.count

        if total_count > 0:
            return Literal(total_weighted / total_count)

        return Literal(float("nan"))

    wrappers["datapoints_average"] = datapoints_average

    # 6. Datapoints Min (using SDK aggregation)
    @safe_sparql_wrapper(default_value=Literal(float("nan")))
    def datapoints_min(
        timeseries_uri: str,
        start: str = "30d-ago",
        end: str = "now",
    ) -> Literal:
        """
        Get minimum value over time range using SDK aggregation.

        Args:
            timeseries_uri: URI containing space and external_id
            start: Start time
            end: End time

        Returns:
            Minimum value as Literal
        """
        timeseries_uri = literal_to_python(timeseries_uri)
        start = _parse_time_param(literal_to_python(start))
        end = _parse_time_param(literal_to_python(end))

        instance_id = parse_instance_id_from_uri(timeseries_uri)

        # Use SDK aggregation with min - use appropriate granularity based on time range
        granularity = "1h"
        if isinstance(start, int) and isinstance(end, int):
            range_ms = end - start
            if range_ms > 86400000:  # More than 1 day
                granularity = "1d"

        result = client.time_series.data.retrieve(
            instance_id=instance_id,
            aggregates=["min"],
            granularity=granularity,
            start=start,
            end=end,
        )

        if result is None or len(result) == 0:
            return Literal(float("nan"))

        # Get minimum across all buckets
        min_value = None
        for datapoint in result:
            if datapoint.min is not None:
                if min_value is None or datapoint.min < min_value:
                    min_value = datapoint.min

        if min_value is not None:
            return Literal(float(min_value))

        return Literal(float("nan"))

    wrappers["datapoints_min"] = datapoints_min

    # 7. Datapoints Max (using SDK aggregation)
    @safe_sparql_wrapper(default_value=Literal(float("nan")))
    def datapoints_max(
        timeseries_uri: str,
        start: str = "30d-ago",
        end: str = "now",
    ) -> Literal:
        """
        Get maximum value over time range using SDK aggregation.

        Args:
            timeseries_uri: URI containing space and external_id
            start: Start time
            end: End time

        Returns:
            Maximum value as Literal
        """
        timeseries_uri = literal_to_python(timeseries_uri)
        start = _parse_time_param(literal_to_python(start))
        end = _parse_time_param(literal_to_python(end))

        instance_id = parse_instance_id_from_uri(timeseries_uri)

        # Use SDK aggregation with max - use appropriate granularity based on time range
        granularity = "1h"
        if isinstance(start, int) and isinstance(end, int):
            range_ms = end - start
            if range_ms > 86400000:  # More than 1 day
                granularity = "1d"

        result = client.time_series.data.retrieve(
            instance_id=instance_id,
            aggregates=["max"],
            granularity=granularity,
            start=start,
            end=end,
        )

        if result is None or len(result) == 0:
            return Literal(float("nan"))

        # Get maximum across all buckets
        max_value = None
        for datapoint in result:
            if datapoint.max is not None:
                if max_value is None or datapoint.max > max_value:
                    max_value = datapoint.max

        if max_value is not None:
            return Literal(float(max_value))

        return Literal(float("nan"))

    wrappers["datapoints_max"] = datapoints_max

    return wrappers
