"""
Cognite Function for partitioned SHACL validation of historic data.

This function queries DMS for a partition of instances (based on range filtering),
validates them with SHACL rules in chunks, and posts results to Records API.

Supports:
- Range-based partitioning (e.g., by date field)
- Cursor-based resumption for fault tolerance
- Chunked processing with progress tracking
- Automatic timeout safety (2-minute margin before workflow timeout)
"""

import time
from typing import Any

from cognite.client.data_classes.data_modeling import ViewId
from cognite.client.data_classes.data_modeling.query import NodeResultSetExpression, Query, Select, SourceSelector
from cognite.client.data_classes.filters import And, Equals, HasData, Range

# Import common utilities
from common.cursor_state import delete_cursor, load_cursor, save_cursor
from common.records_api import post_validation_results_to_records
from common.shacl_utils import load_shacl_from_file


def handle_partitioned_validation(data: dict[str, Any], client) -> dict[str, Any]:
    """
    Validate a partition of historic instances with SHACL rules.

    Args:
        data: Input data containing:
            - orchestration_id: Orchestration run identifier (links to OrchestrationState)
            - partition_id: Unique ID for this partition (e.g., "0", "1", "2")
            - partition_range: Dict with min_value and max_value for the partition_field
            - partition_field: Field to filter on (e.g., "scheduledStartTime")

            - view_space: Space of the view to query
            - view_external_id: External ID of the view
            - view_version: Version of the view
            - instance_space: Space where instances live

            - datamodel_space: Space of the data model
            - datamodel_external_id: External ID of the data model
            - datamodel_version: Version of the data model

            - shacl_file_external_id: External ID of the SHACL rules file
            - rule_set_id: Rule set identifier
            - rule_set_version: Rule set version
            - stream_id: Stream ID for Records API

            - auto_load_depth: (optional) Depth for auto-loading references (default: 2)
            - chunk_size: (optional) Number of instances per chunk (default: 200)
            - max_runtime_seconds: (optional) Maximum runtime before timeout (default: 480 = 8 minutes)

        client: Cognite client instance

    Returns:
        Dict with:
            - status: "completed" or "partial" (if resumed later)
            - partition_id: The partition ID
            - instances_validated: Total instances validated in this run
            - total_violations: Total violations found
            - records_posted: Total records posted
            - cursor_saved: Whether a cursor was saved for resumption
    """
    from cognite_data_quality._neat import NeatSession

    print("=" * 80)
    print("PARTITIONED HISTORIC VALIDATION")
    print("=" * 80)

    # Extract parameters
    orchestration_id = data.get("orchestration_id", "unknown")
    partition_id = data.get("partition_id", "unknown")
    partition_range = data.get("partition_range", {})
    partition_field = data.get("partition_field")

    view_space = data.get("view_space")
    view_external_id = data.get("view_external_id")
    view_version = data.get("view_version")
    instance_space = data.get("instance_space")

    datamodel_space = data.get("datamodel_space")
    datamodel_external_id = data.get("datamodel_external_id")
    datamodel_version = data.get("datamodel_version")

    shacl_file_external_id = data.get("shacl_file_external_id")
    rule_set_id = data.get("rule_set_id")
    rule_set_version = data.get("rule_set_version")
    stream_id = data.get("stream_id")

    # Data quality space for state containers (default: "dataQuality")
    data_quality_space = data.get("data_quality_space", "dataQuality")

    # Records API configuration
    records_space = data.get("records_space", "dataQuality")
    records_container = data.get("records_container", "DataQualityValidationRecord")

    auto_load_depth = data.get("auto_load_depth", 2)
    chunk_size = data.get("chunk_size", 200)
    max_runtime_seconds = data.get(
        "max_runtime_seconds", 480
    )  # 8 minutes (2-min safety margin before 10-min workflow timeout)

    # Use passed job_run_id (shared across all partitions) or generate one if not provided
    job_run_id = data.get("job_run_id", f"historic_{partition_id}_{int(time.time())}")

    print(f"Partition ID: {partition_id}")
    print(f"Job Run ID: {job_run_id}")
    print(f"Partition range: {partition_range}")
    print(f"Partition field: {partition_field}")
    print(f"View: {view_space}/{view_external_id}/{view_version}")
    print(f"Instance space: {instance_space}")
    print(f"Data model: {datamodel_space}/{datamodel_external_id}/{datamodel_version}")
    print(f"Chunk size: {chunk_size}")
    print(f"Max runtime: {max_runtime_seconds}s")

    start_time = time.time()

    # Load SHACL rules
    try:
        print(f"Loading SHACL rules from: {shacl_file_external_id}")
        shacl_rules = load_shacl_from_file(client, shacl_file_external_id)
        print(f"  Loaded {len(shacl_rules)} characters")
    except Exception as e:
        return {"status": "error", "error": f"Failed to load SHACL rules: {e!s}", "partition_id": partition_id}

    # Build dimension lookup from SHACL rules (parsed once; O(1) lookup per violation)
    from rdflib import BNode as _BNode
    from rdflib import Graph as _Graph
    from rdflib import Namespace as _Namespace

    _shape_to_dim: dict[str, str] = {}
    _path_to_dim: dict[str, str] = {}
    _path_to_shape: dict[str, str] = {}
    _shape_to_targetclass: dict[str, str] = {}
    try:
        _shacl_graph = _Graph()
        _shacl_graph.parse(data=shacl_rules, format="turtle")
        _dq = _Namespace("http://example.org/dataquality/")
        _sh_ns = _Namespace("http://www.w3.org/ns/shacl#")
        for _shape, _dim in _shacl_graph.subject_objects(predicate=_dq.dimension):
            if isinstance(_shape, _BNode):
                for _path in _shacl_graph.objects(subject=_shape, predicate=_sh_ns.path):
                    _path_to_dim[str(_path)] = str(_dim)
            else:
                _shape_to_dim[str(_shape)] = str(_dim)
        # Build path → parent named shape mapping (stable key across graph re-parsings)
        for _parent, _prop_bnode in _shacl_graph.subject_objects(predicate=_sh_ns.property):
            if not isinstance(_parent, _BNode):
                _prop_path = _shacl_graph.value(_prop_bnode, _sh_ns.path)
                if _prop_path and not isinstance(_prop_path, _BNode):
                    _path_to_shape[str(_prop_path)] = str(_parent)
        # Build shape → targetClass local name mapping
        for _shape, _tc in _shacl_graph.subject_objects(predicate=_sh_ns.targetClass):
            if not isinstance(_shape, _BNode):
                _tc_str = str(_tc)
                _shape_to_targetclass[str(_shape)] = (
                    _tc_str.split("#")[-1] if "#" in _tc_str else _tc_str.split("/")[-1]
                )
        print(f"  Dimension lookup: {len(_shape_to_dim)} named shapes, {len(_path_to_dim)} property paths")
    except Exception as _dim_err:
        print(f"  Warning: Could not build dimension lookup: {_dim_err}")

    # Load cursor if exists (resuming from previous run)
    print("Looking for saved cursor:")
    print(f"  Orchestration ID: {orchestration_id}")
    print(f"  Partition ID: {partition_id}")
    print(f"  Expected external_id: partition_state_orch_{orchestration_id}_p{partition_id}")

    cursor_state = load_cursor(client, orchestration_id, partition_id, data_quality_space)
    start_cursor = cursor_state["cursor"] if cursor_state else None
    total_processed = cursor_state["processed_count"] if cursor_state else 0

    if start_cursor:
        print(f"✓ Resuming from cursor (already processed: {total_processed})")
        print(f"  Cursor: {start_cursor[:80]}...")
    else:
        print("✗ Starting fresh (no saved cursor found)")

    # Create ViewId
    view_id = ViewId(view_space, view_external_id, view_version)

    # Build DMS filter with range
    filters = [HasData(views=[view_id]), Equals(["node", "space"], instance_space)]

    if partition_range:
        min_val = partition_range.get("min_value")
        max_val = partition_range.get("max_value")

        if min_val is not None and max_val is not None:
            # Check if partition_field is a system property
            system_properties = ["lastUpdatedTime", "createdTime", "externalId", "space", "type"]

            if partition_field in system_properties:
                # System properties use ["node", property_name] path
                filters.append(
                    Range(
                        property=["node", partition_field],
                        gte=min_val,
                        lt=max_val,  # Exclusive upper bound to avoid overlap
                    )
                )
                print(f"Using range filter on system property: {partition_field} in [{min_val}, {max_val})")
                print(f'  System property path: ["node", "{partition_field}"]')
            else:
                # View properties need to be resolved to container properties
                # Range filters require container properties, not view properties
                try:
                    view_def_list = client.data_modeling.views.retrieve((view_space, view_external_id, view_version))
                    view_def = view_def_list[0] if view_def_list and len(view_def_list) > 0 else None

                    if view_def and view_def.properties and partition_field in view_def.properties:
                        prop_def = view_def.properties[partition_field]
                        if hasattr(prop_def, "container"):
                            container_ref = prop_def.container
                            container_property = prop_def.container_property_identifier or partition_field

                            filters.append(
                                Range(
                                    property=[container_ref.space, container_ref.external_id, container_property],
                                    gte=min_val,
                                    lt=max_val,  # Exclusive upper bound to avoid overlap
                                )
                            )
                            print(f"Using range filter: {partition_field} in [{min_val}, {max_val})")
                            print(
                                f"  Container path: [{container_ref.space},"
                                f" {container_ref.external_id}, {container_property}]"
                            )
                        else:
                            print(f"  Warning: No container reference for {partition_field}, skipping range filter")
                    else:
                        print(f"  Warning: Property {partition_field} not found in view, skipping range filter")
                except Exception as e:
                    print(f"  Warning: Could not resolve container for {partition_field}: {e}")
                    print("  Skipping range filter - will process all instances")

    dms_filter = And(*filters)

    # Initialize NEAT session
    neat = NeatSession(client, verbose=False)

    # Track totals
    total_instances_validated = 0
    total_violations_found = 0
    total_records_posted = 0

    # Process in chunks with cursor
    current_cursor = start_cursor
    chunk_num = 0

    # Adaptive sync limit sequence — mirrors timestamp collection logic.
    # Index persists across chunks so we don't restart at chunk_size after a timeout.
    # Build sequence starting at chunk_size, then step through standard fallbacks below it.
    _base_limits = [1000, 500, 200, 100, 10]
    SYNC_LIMITS = [chunk_size] + [v for v in _base_limits if v < chunk_size]
    sync_limit_idx = 0

    try:
        while True:
            # Check timeout (5-minute safety margin)
            elapsed = time.time() - start_time
            if elapsed > max_runtime_seconds:
                print(f"Approaching timeout ({elapsed:.1f}s), saving cursor and exiting...")
                if current_cursor:
                    save_cursor(
                        client,
                        orchestration_id,
                        partition_id,
                        current_cursor,
                        total_processed + total_instances_validated,
                        data_quality_space,
                    )

                return {
                    "status": "partial",
                    "partition_id": partition_id,
                    "instances_validated": total_instances_validated,
                    "cumulative_validated": total_processed + total_instances_validated,
                    "total_violations": total_violations_found,
                    "records_posted": total_records_posted,
                    "cursor_saved": True,
                    "elapsed_seconds": elapsed,
                }

            chunk_num += 1
            print(f"--- Chunk {chunk_num} (elapsed: {elapsed:.1f}s) ---")

            # Query next chunk from DMS using sync with adaptive limit (steps down on timeout)
            result = None
            retries_at_limit = 0
            while True:
                sync_limit = SYNC_LIMITS[sync_limit_idx]
                try:
                    query = Query(
                        with_={
                            "nodes": NodeResultSetExpression(
                                filter=dms_filter,
                                limit=sync_limit,
                                sync_mode="two_phase",
                            )
                        },
                        select={"nodes": Select(sources=[SourceSelector(source=view_id, properties=[])])},
                        cursors={"nodes": current_cursor} if current_cursor else None,
                    )
                    result = client.data_modeling.instances.sync(query=query)
                    break  # success
                except Exception as e:
                    error_msg = str(e).lower()
                    if "timeout" in error_msg or "timed out" in error_msg:
                        retries_at_limit += 1
                        if retries_at_limit < 2:
                            print(
                                f"  Sync timeout (attempt {retries_at_limit}/2),"
                                f" retrying with limit={sync_limit} in 30s..."
                            )
                            time.sleep(30)
                        else:
                            retries_at_limit = 0
                            sync_limit_idx += 1
                            if sync_limit_idx >= len(SYNC_LIMITS):
                                print("  All sync limits exhausted — saving cursor and aborting partition")
                                if current_cursor:
                                    save_cursor(
                                        client,
                                        orchestration_id,
                                        partition_id,
                                        current_cursor,
                                        total_processed + total_instances_validated,
                                        data_quality_space,
                                    )
                                return {
                                    "status": "sync_limits_exhausted",
                                    "partition_id": partition_id,
                                    "instances_validated": total_instances_validated,
                                    "cumulative_validated": total_processed + total_instances_validated,
                                    "total_violations": total_violations_found,
                                    "records_posted": total_records_posted,
                                    "cursor_saved": bool(current_cursor),
                                    "message": "All sync limits exhausted, cursor saved for orchestrator retry",
                                }
                            print(
                                f"  Sync timeout (2/2), reducing limit:"
                                f" {sync_limit} → {SYNC_LIMITS[sync_limit_idx]}, waiting 30s..."
                            )
                            time.sleep(30)
                    else:
                        raise

            # Extract instances from sync result and convert to dicts
            instances = []
            if hasattr(result, "data") and "nodes" in result.data:
                node_objects = result.data["nodes"]
                # Convert SDK Node objects to dictionaries with camelCase keys (externalId, space)
                instances = [node.dump(camel_case=True) if hasattr(node, "dump") else node for node in node_objects]

            # Update cursor for next iteration
            new_cursor = None
            if hasattr(result, "cursors"):
                if hasattr(result.cursors, "__dict__"):
                    cursors_dict = result.cursors.__dict__
                    new_cursor = cursors_dict.get("nodes")
                elif isinstance(result.cursors, dict):
                    new_cursor = result.cursors.get("nodes")

            has_more = new_cursor is not None and new_cursor != current_cursor
            current_cursor = new_cursor

            print(f"  Fetched {len(instances)} instances (has_more: {has_more})")

            if not instances:
                print("  No more instances to process")
                break

            # Validate chunk
            _conforms, report_graph, _report_text = neat.validate_instances.with_shacl(
                instances=instances,
                shacl_rules=shacl_rules,
                datamodel_space=datamodel_space,
                datamodel_external_id=datamodel_external_id,
                datamodel_version=datamodel_version,
                auto_load_depth=auto_load_depth,
                verbose=False,
            )

            # Extract violations
            violations = []
            if report_graph:
                from rdflib import Graph, Namespace

                if not isinstance(report_graph, Graph):
                    # pyshacl caught a ValidationFailure internally (e.g. a SPARQL constraint
                    # crashed on unexpected data) and returned the exception object as report_graph.
                    # Create one synthetic violation per instance so Records API reflects that
                    # these instances could not be validated, rather than silently passing them.
                    engine_error = str(report_graph)
                    print(f"  WARNING: SHACL engine failure for this chunk: {engine_error}")
                    for inst in instances:
                        inst_space = inst.get("space", "")
                        inst_ext_id = inst.get("externalId", "")
                        violations.append(
                            {
                                "focusNode": f"http://purl.org/cognite/{inst_space}/{inst_ext_id}",
                                "sourceConstraintComponent": "SHACLEngineFailure",
                                "resultMessage": f"SHACL engine failure prevented validation: {engine_error}",
                                "resultSeverity": "Violation",
                            }
                        )
                else:
                    SH = Namespace("http://www.w3.org/ns/shacl#")

                    for result_node in report_graph.subjects(predicate=None, object=SH.ValidationResult):
                        violation = {}
                        for pred, obj in report_graph.predicate_objects(subject=result_node):
                            pred_name = str(pred).split("#")[-1]
                            if pred_name == "sourceShape" and isinstance(obj, _BNode):
                                pass  # blank node IDs don't survive re-parsing; resolve via resultPath below
                            else:
                                violation[pred_name] = str(obj)
                        # Resolve missing sourceShape via resultPath → parent named shape
                        if "sourceShape" not in violation:
                            _rp = violation.get("resultPath")
                            if _rp and _rp in _path_to_shape:
                                violation["sourceShape"] = _path_to_shape[_rp]
                        violations.append(violation)

            print(f"  Validated {len(instances)} instances, found {len(violations)} violations")

            # Attach dimension and targetClass to each violation via dict lookup
            for v in violations:
                source_shape = v.get("sourceShape", "")
                result_path = v.get("resultPath", "")
                if source_shape and source_shape in _shape_to_dim:
                    v["dimension"] = _shape_to_dim[source_shape]
                elif result_path and result_path in _path_to_dim:
                    v["dimension"] = _path_to_dim[result_path]
                if source_shape and source_shape in _shape_to_targetclass:
                    v["targetClass"] = _shape_to_targetclass[source_shape]

            # Post to Records API
            namespace_base = "http://purl.org/cognite/"
            posted_count, record_errors = post_validation_results_to_records(
                client=client,
                all_instances=instances,
                violations=violations,
                job_run_id=job_run_id,
                stream_id=stream_id,
                rule_set_id=rule_set_id,
                rule_set_version=rule_set_version,
                namespace_base=namespace_base,
                records_space=records_space,
                records_container=records_container,
            )

            print(f"  Posted {posted_count} records")
            if record_errors:
                print(f"  WARNING: {len(record_errors)} record(s) failed to post:")
                for err in record_errors[:5]:  # Show first 5 errors
                    print(f"    - {err['instance']}: {err['error']}")
                if len(record_errors) > 5:
                    print(f"    ... and {len(record_errors) - 5} more")

                # Fail partition if posting failure rate exceeds threshold
                failure_rate = len(record_errors) / len(instances)
                failure_threshold = 0.1  # 10% threshold
                if failure_rate > failure_threshold:
                    raise RuntimeError(
                        f"Failed to post {len(record_errors)}/{len(instances)} records "
                        f"({failure_rate:.1%} failure rate exceeds {failure_threshold:.1%} threshold). "
                        f"Validation results cannot be persisted reliably. "
                        f"This is likely a transient API issue - the partition will be retried."
                    )

            # Update totals
            total_instances_validated += len(instances)
            total_violations_found += len(violations)
            total_records_posted += posted_count

            # Save cursor after successful chunk (for resumption)
            # If cursor save fails, we must fail the partition to avoid losing progress
            if has_more and current_cursor:
                try:
                    save_cursor(
                        client,
                        orchestration_id,
                        partition_id,
                        current_cursor,
                        total_processed + total_instances_validated,
                        data_quality_space,
                    )
                except Exception as cursor_error:
                    # This is critical - if we can't save progress, we can't safely continue
                    raise RuntimeError(
                        f"Failed to save cursor after processing {total_instances_validated} instances. "
                        f"Cannot continue without progress tracking. Error: {cursor_error}"
                    ) from cursor_error

            # If no more data, exit loop
            if not has_more:
                print("  Reached end of partition")
                break

        # Verify all records were posted successfully before completing
        if total_records_posted < total_instances_validated:
            # Some records failed to post - don't delete cursor, fail the partition
            raise RuntimeError(
                f"Only {total_records_posted}/{total_instances_validated} records posted successfully. "
                f"Cannot mark partition as completed with missing records. "
                f"Cursor state preserved for retry."
            )

        # All records posted successfully - delete cursor state
        delete_cursor(client, orchestration_id, partition_id, data_quality_space)

        elapsed = time.time() - start_time
        print("=" * 80)
        print(f"PARTITION {partition_id} COMPLETED")
        cumulative_total = total_processed + total_instances_validated

        print("=" * 80)
        print(f"Instances validated (this run): {total_instances_validated}")
        print(f"Cumulative validated: {cumulative_total}")
        print(f"Total violations: {total_violations_found}")
        print(f"Records posted: {total_records_posted}")
        print(f"Elapsed time: {elapsed:.1f}s")
        print("=" * 80)

        return {
            "status": "completed",
            "partition_id": partition_id,
            "instances_validated": total_instances_validated,
            "cumulative_validated": cumulative_total,
            "total_violations": total_violations_found,
            "records_posted": total_records_posted,
            "cursor_saved": False,
            "elapsed_seconds": elapsed,
        }

    except Exception as e:
        import traceback

        error_msg = f"Error in partition {partition_id}: {e!s}\n{traceback.format_exc()}"
        print(f"{'=' * 80}")
        print(f"PARTITION {partition_id} FAILED")
        print(f"{'=' * 80}")
        print(error_msg)
        print(f"{'=' * 80}")

        # Save cursor on error for retry - fail if this doesn't work
        if current_cursor:
            save_cursor(
                client,
                orchestration_id,
                partition_id,
                current_cursor,
                total_processed + total_instances_validated,
                data_quality_space,
            )
            print("Cursor saved for retry")

        # Return error status - orchestrator should detect this
        return {
            "status": "error",
            "error": error_msg,
            "partition_id": partition_id,
            "instances_validated": total_instances_validated,
            "cumulative_validated": total_processed + total_instances_validated,
            "total_violations": total_violations_found,
            "records_posted": total_records_posted,
            "cursor_saved": bool(current_cursor),
        }
