"""
Cursor state management for SHACL validation functions.

This module consolidates cursor state management for two validation types:
1. Partitioned validation: Tracks progress through partitioned instance validation
2. Time series validation: Tracks validation state per time series per rule set

Functions prefixed with 'partitioned_' are for partitioned validation.
Functions prefixed with 'timeseries_' are for time series validation.
"""

from __future__ import annotations

import datetime
import time
from typing import Any

from cognite.client import CogniteClient

# Constants for time series validation cursor storage
CURSOR_CONTAINER_ID = "TimeSeriesValidationCursor"
CURSOR_VIEW_ID = "TimeSeriesValidationCursorView"
CURSOR_VIEW_VERSION = "v1"


# ============================================================================
# Partitioned Validation Cursor Functions
# ============================================================================


def load_cursor(
    client, orchestration_id: str, partition_id: str, data_quality_space: str = "dataQuality"
) -> dict | None:
    """
    Load the saved cursor for this partition from the state container.

    For partitioned validation: Tracks progress through a single partition.

    Returns None if no cursor exists (first run or completed partition).
    """
    try:
        external_id = f"partition_state_orch_{orchestration_id}_p{partition_id}"

        # Query for the specific node by external_id with sources to get properties
        response = client.post(
            url=f"/api/v1/projects/{client.config.project}/models/instances/list",
            json={
                "instanceType": "node",
                "sources": [
                    {
                        "source": {
                            "type": "view",
                            "space": data_quality_space,
                            "externalId": "FunctionValidationStateView",
                            "version": "1",
                        }
                    }
                ],
                "filter": {
                    "and": [
                        {"equals": {"property": ["node", "space"], "value": data_quality_space}},
                        {"equals": {"property": ["node", "externalId"], "value": external_id}},
                    ]
                },
                "limit": 1,
            },
        )

        result_data = response.json()
        items = result_data.get("items", [])

        if not items:
            return None

        # Extract properties from the FunctionValidationState view
        item = items[0]
        print(f"  Found state record: {item.get('externalId')}, version: {item.get('version')}")

        # Properties structure: {space/externalId/version: {container: {property: value}}}
        properties = item.get("properties", {})

        # Find the properties in any source (usually first one)
        for _source_key, containers in properties.items():
            for _container_key, props in containers.items():
                cursor = props.get("syncCursor")
                processed_count = props.get("processedCount", 0)

                print(f"  Loaded properties: syncCursor={bool(cursor)}, processedCount={processed_count}")

                if cursor:
                    print(
                        f"  ✓ Loaded cursor for partition {partition_id}: {cursor[:50]}..."
                        f" (already processed: {processed_count})"
                    )
                    return {"cursor": cursor, "processed_count": processed_count}

        # No cursor found
        print("  ✗ No cursor found in state record (starting fresh)")
        return None

    except Exception as e:
        print(f"Failed to load cursor: {e}")
        return None


def save_cursor(
    client,
    orchestration_id: str,
    partition_id: str,
    cursor: str,
    processed_count: int,
    data_quality_space: str = "dataQuality",
):
    """
    Save the current cursor for this partition to the state container.

    For partitioned validation: Saves progress through a single partition.

    Raises exception if save fails - this is critical for resumption.
    """
    try:
        from cognite.client.data_classes.data_modeling.ids import ContainerId
        from cognite.client.data_classes.data_modeling.instances import NodeApply, NodeOrEdgeData

        external_id = f"partition_state_orch_{orchestration_id}_p{partition_id}"

        print(f"Saving cursor for partition {partition_id}...")
        print(f"  External ID: {external_id}")
        print(f"  Processed: {processed_count}")
        print(f"  Cursor: {cursor[:50] if cursor else 'None'}...")

        # Create source with proper SDK structure
        source = NodeOrEdgeData(
            source=ContainerId(data_quality_space, "FunctionValidationState"),
            properties={
                "orchestration_id": orchestration_id,
                "partitionId": str(partition_id),  # Must be string for container schema
                "syncCursor": cursor,
                "processedCount": processed_count,
                "lastUpdated": int(time.time() * 1000),  # milliseconds
            },
        )

        node = NodeApply(space=data_quality_space, external_id=external_id, sources=[source])

        # Upsert the state record
        result = client.data_modeling.instances.apply(nodes=[node])

        print(f"✓ Saved cursor for partition {partition_id}")
        print(f"  Result: {result}")

        # Verify the write by reading it back using the view
        try:
            from cognite.client.data_classes.data_modeling.ids import ViewId

            verify = client.data_modeling.instances.retrieve(
                nodes=[(data_quality_space, external_id)],
                sources=[ViewId(data_quality_space, "FunctionValidationStateView", "1")],
            )
            if verify and verify.nodes:
                print("✓ Verified: cursor state written successfully")
            else:
                print("⚠ Warning: Could not verify cursor state after write")
        except Exception as verify_error:
            print(f"⚠ Warning: Could not verify cursor write: {verify_error}")

    except Exception as e:
        error_msg = f"CRITICAL ERROR saving cursor for partition {partition_id}: {e}"
        print(error_msg)
        print(f"  Error type: {type(e).__name__}")
        print(f"  Error details: {e!s}")
        raise RuntimeError(error_msg) from e


def load_sync_cursor(
    client,
    view_space: str,
    view_external_id: str,
    instance_space: str,
    job_run_id: str,
    data_quality_space: str = "dataQuality",
) -> dict | None:
    """
    Load the saved sync cursor for a view/instance_space pair (instance_sync_cursor mode).

    External ID: sync_cursor_{job_run_id}_{view_space}_{view_external_id}_{instance_space}

    Returns None if no cursor exists (first run).
    """
    try:
        external_id = f"sync_cursor_{job_run_id}_{view_space}_{view_external_id}_{instance_space}"

        response = client.post(
            url=f"/api/v1/projects/{client.config.project}/models/instances/list",
            json={
                "instanceType": "node",
                "sources": [
                    {
                        "source": {
                            "type": "view",
                            "space": data_quality_space,
                            "externalId": "FunctionValidationStateView",
                            "version": "1",
                        }
                    }
                ],
                "filter": {
                    "and": [
                        {"equals": {"property": ["node", "space"], "value": data_quality_space}},
                        {"equals": {"property": ["node", "externalId"], "value": external_id}},
                    ]
                },
                "limit": 1,
            },
        )

        result_data = response.json()
        items = result_data.get("items", [])

        if not items:
            return None

        item = items[0]
        properties = item.get("properties", {})

        for _source_key, containers in properties.items():
            for _container_key, props in containers.items():
                cursor = props.get("syncCursor")
                processed_count = props.get("processedCount", 0)

                if cursor:
                    print(f"  ✓ Loaded sync cursor for {view_external_id}/{instance_space}: {cursor[:50]}...")
                    return {"cursor": cursor, "processed_count": processed_count}

        return None

    except Exception as e:
        print(f"Failed to load sync cursor: {e}")
        return None


def save_sync_cursor(
    client,
    view_space: str,
    view_external_id: str,
    instance_space: str,
    cursor: str,
    processed_count: int,
    job_run_id: str,
    data_quality_space: str = "dataQuality",
):
    """
    Save the sync cursor for a view/instance_space pair (instance_sync_cursor mode).

    Raises exception if save fails — this is critical for resumption.
    """
    try:
        from cognite.client.data_classes.data_modeling.ids import ContainerId
        from cognite.client.data_classes.data_modeling.instances import NodeApply, NodeOrEdgeData

        external_id = f"sync_cursor_{job_run_id}_{view_space}_{view_external_id}_{instance_space}"

        source = NodeOrEdgeData(
            source=ContainerId(data_quality_space, "FunctionValidationState"),
            properties={
                "orchestration_id": f"{view_space}/{view_external_id}",
                "partitionId": "sync",
                "syncCursor": cursor,
                "processedCount": processed_count,
                "lastUpdated": int(time.time() * 1000),
            },
        )

        node = NodeApply(space=data_quality_space, external_id=external_id, sources=[source])
        client.data_modeling.instances.apply(nodes=[node])
        print(f"✓ Saved sync cursor for {view_external_id}/{instance_space} (processed: {processed_count})")

    except Exception as e:
        error_msg = f"CRITICAL ERROR saving sync cursor for {view_external_id}/{instance_space}: {e}"
        print(error_msg)
        raise RuntimeError(error_msg) from e


def delete_cursor(client, orchestration_id: str, partition_id: str, data_quality_space: str = "dataQuality"):
    """
    Delete the cursor state after successful completion.

    For partitioned validation: Removes cursor after partition is complete.
    """
    try:
        external_id = f"partition_state_orch_{orchestration_id}_p{partition_id}"
        client.data_modeling.instances.delete(nodes=[(data_quality_space, external_id)])
        print(f"Deleted cursor state for partition {partition_id}")
    except Exception as e:
        print(f"Failed to delete cursor state: {e}")


# ============================================================================
# Time Series Validation Cursor Functions
# ============================================================================


def ensure_cursor_container(client: CogniteClient, data_quality_space: str = "dataQuality") -> None:
    """
    Create the cursor storage container, view, and data model if they don't exist.

    For time series validation: Sets up DMS structures for cursor storage.
    """
    from cognite.client.data_classes.data_modeling import (
        ContainerApply,
        ContainerId,
        ContainerProperty,
        DataModelApply,
        MappedPropertyApply,
        SpaceApply,
        Text,
        ViewApply,
        ViewId,
    )

    space = data_quality_space
    container_id = CURSOR_CONTAINER_ID
    view_id = f"{container_id}View"
    model_id = "TimeSeriesValidationState"
    version = "v1"

    print(f"  Setting up cursor storage: {space}/{container_id}")

    # 1. Ensure space exists
    try:
        client.data_modeling.spaces.apply(SpaceApply(space=space))
        print(f"    Space '{space}' ready")
    except Exception as e:
        print(f"    Space '{space}' error: {e}")

    # 2. Create container
    container = ContainerApply(
        space=space,
        external_id=container_id,
        name="Time Series Validation Cursors",
        description="Stores validation state per time series per rule set",
        properties={
            "timeseriesId": ContainerProperty(type=Text(), nullable=False),
            "ruleSetId": ContainerProperty(type=Text(), nullable=False),
            "lastValidatedTimestamp": ContainerProperty(type=Text(), nullable=True),
            "lastStatus": ContainerProperty(type=Text(), nullable=True),
            "backfillProgress": ContainerProperty(type=Text(), nullable=True),
            "subscriptionCursor": ContainerProperty(type=Text(), nullable=True),
        },
    )

    try:
        client.data_modeling.containers.apply(container)
        print(f"    Container '{container_id}' ready")
    except Exception as e:
        print(f"    Container '{container_id}' error: {e}")

    # 3. Create view
    view = ViewApply(
        space=space,
        external_id=view_id,
        version=version,
        name="Time Series Validation Cursors",
        description="View for validation state tracking",
        properties={
            "timeseriesId": MappedPropertyApply(
                container=ContainerId(space=space, external_id=container_id),
                container_property_identifier="timeseriesId",
            ),
            "ruleSetId": MappedPropertyApply(
                container=ContainerId(space=space, external_id=container_id),
                container_property_identifier="ruleSetId",
            ),
            "lastValidatedTimestamp": MappedPropertyApply(
                container=ContainerId(space=space, external_id=container_id),
                container_property_identifier="lastValidatedTimestamp",
            ),
            "lastStatus": MappedPropertyApply(
                container=ContainerId(space=space, external_id=container_id),
                container_property_identifier="lastStatus",
            ),
            "backfillProgress": MappedPropertyApply(
                container=ContainerId(space=space, external_id=container_id),
                container_property_identifier="backfillProgress",
            ),
            "subscriptionCursor": MappedPropertyApply(
                container=ContainerId(space=space, external_id=container_id),
                container_property_identifier="subscriptionCursor",
            ),
        },
    )

    try:
        client.data_modeling.views.apply(view)
        print(f"    View '{view_id}/{version}' ready")
    except Exception as e:
        print(f"    View '{view_id}/{version}' error: {e}")

    # 4. Create data model
    data_model = DataModelApply(
        space=space,
        external_id=model_id,
        version=version,
        name="Time Series Validation State",
        description="Data model for tracking validation state and cursors",
        views=[ViewId(space=space, external_id=view_id, version=version)],
    )

    try:
        client.data_modeling.data_models.apply(data_model)
        print(f"    Data model '{model_id}/{version}' ready")
    except Exception as e:
        print(f"    Data model '{model_id}/{version}' error: {e}")


def load_cursor_state(
    client: CogniteClient,
    timeseries_ids: list[str],
    rule_set_id: str,
    data_quality_space: str = "dataQuality",
) -> dict[str, dict[str, Any]]:
    """
    Load cursor state for given time series from DMS.

    For time series validation: Loads validation state per time series per rule set.

    Cursors are stored with external_id = "{rule_set_id}/{timeseries_id}"
    Returns state keyed by timeseries_id for easier lookup.
    """
    if not timeseries_ids:
        return {}

    try:
        from cognite.client.data_classes.data_modeling import NodeId, ViewId

        # Cursor IDs include rule_set_id prefix
        cursor_node_ids = [
            NodeId(space=data_quality_space, external_id=f"{rule_set_id}/{ts_id}") for ts_id in timeseries_ids
        ]

        # Must specify view as source to get properties
        view_id = ViewId(space=data_quality_space, external_id=CURSOR_VIEW_ID, version=CURSOR_VIEW_VERSION)

        # Use retrieve with sources to get properties
        result = client.data_modeling.instances.retrieve(
            nodes=cursor_node_ids,
            sources=[view_id],
        )

        state = {}
        # Properties are keyed by: space -> "viewId/version" -> property
        view_key = f"{CURSOR_VIEW_ID}/{CURSOR_VIEW_VERSION}"
        for node in result.nodes:
            # Use .dump() to get raw dict - .get() expects ViewId not string
            props_dict = node.properties.dump()
            props = props_dict.get(data_quality_space, {}).get(view_key, {})
            # Extract timeseries_id from properties or external_id
            ts_id = props.get("timeseriesId")
            if not ts_id and "/" in node.external_id:
                # Fallback: extract from external_id (rule_set_id/ts_id)
                ts_id = node.external_id.split("/", 1)[-1]
            if ts_id:
                state[ts_id] = {
                    "lastValidatedTimestamp": props.get("lastValidatedTimestamp"),
                    "lastStatus": props.get("lastStatus"),
                    "backfillProgress": props.get("backfillProgress"),
                    "subscriptionCursor": props.get("subscriptionCursor"),
                }
        return state

    except Exception as e:
        import logging

        logger = logging.getLogger(__name__)
        logger.warning(f"Failed to load cursor state: {e}")
        return {}


def save_cursor_state(
    client: CogniteClient,
    timeseries_id: str,
    rule_set_id: str,
    status: str,
    data_quality_space: str = "dataQuality",
    backfill_progress: str | None = None,
    subscription_cursor: str | None = None,
) -> None:
    """
    Save cursor state for a time series + rule set combination to DMS.

    For time series validation: Saves validation state per time series per rule set.
    """
    from cognite.client.data_classes.data_modeling import ContainerId, NodeApply, NodeOrEdgeData

    now = datetime.datetime.now(tz=datetime.timezone.utc).isoformat()
    # Unique ID per time series per rule set
    cursor_id = f"{rule_set_id}/{timeseries_id}"

    node = NodeApply(
        space=data_quality_space,
        external_id=cursor_id,
        sources=[
            NodeOrEdgeData(
                source=ContainerId(space=data_quality_space, external_id=CURSOR_CONTAINER_ID),
                properties={
                    "timeseriesId": timeseries_id,
                    "ruleSetId": rule_set_id,
                    "lastValidatedTimestamp": now,
                    "lastStatus": status,
                    "backfillProgress": backfill_progress,
                    "subscriptionCursor": subscription_cursor,
                },
            )
        ],
    )

    try:
        client.data_modeling.instances.apply(node)
    except Exception as e:
        print(f"  ERROR saving cursor for {cursor_id}: {e}")
