"""
Cognite Function for SHACL validation of RAW table rows.

This function validates RAW table rows against SHACL rules, supporting both
incremental validation (timestamp-based) and historic validation (cursor-based).

Validation results are posted to CDF Records API for data quality tracking.
"""

import time
import time as time_module
from typing import Any

from common.records_api import post_validation_results_to_records
from common.shacl_utils import load_shacl_from_file

from cognite_data_quality._neat import NeatSession


def handle_raw_validation_incremental(data: dict[str, Any], client) -> dict[str, Any]:
    """
    Validate RAW table rows incrementally with SHACL.

    Used for scheduled workflows that run frequently (e.g., every 15 minutes).
    Validates only rows updated since last run based on lastUpdatedTime.

    Args:
        data: Input data containing:
            - db_name: RAW database name
            - table_name: RAW table name
            - shacl_rules: SHACL rules as Turtle string (use this OR shacl_rules_file_external_id)
            - shacl_rules_file_external_id: External ID of a CDF File containing SHACL rules
            - batch_size: (optional) Maximum rows to validate (default: 10000)
            - verbose: (optional) Print progress messages (default: True)
            - job_run_id: (optional) Unique job run ID for records
            - records_config: Config for Records API posting:
                - stream_id: Stream ID for records
                - rule_set_id: Rule set identifier
                - rule_set_version: Rule set version
                - data_domain_external_id: Data domain for records
                - records_space: (optional) Space for records (default: "dataQuality")
                - records_container: (optional) Container name (default: "DataQualityValidationRecord")
        client: Cognite client instance

    Returns:
        Dict with validation results:
            - conforms: Boolean indicating if validation passed
            - rows_validated: Number of rows validated
            - violations_found: Number of violations found
            - last_processed_timestamp: Timestamp for next incremental run
            - records_posted: Number of records posted to Records API
    """
    print("=" * 80)
    print("RAW TABLE SHACL VALIDATION (INCREMENTAL)")
    print("=" * 80)

    # Extract parameters
    db_name = data.get("db_name")
    table_name = data.get("table_name")
    shacl_rules = data.get("shacl_rules")
    shacl_rules_file_external_id = data.get("shacl_rules_file_external_id")
    batch_size = data.get("batch_size", 10000)
    verbose = data.get("verbose", True)
    data_quality_space = data.get("data_quality_space", "dataQuality")

    # Records API parameters
    job_run_id = data.get("job_run_id") or f"job_{int(time.time())}"
    records_config = data.get("records_config", {})

    # Validate required parameters
    if not db_name or not table_name:
        return {
            "conforms": False,
            "error": "Missing required parameters: db_name, table_name",
            "violations_found": 0,
        }

    # Load SHACL rules from CDF File if external ID is provided
    if shacl_rules_file_external_id and not shacl_rules:
        try:
            print(f"Loading SHACL rules from CDF File: {shacl_rules_file_external_id}")
            shacl_rules = load_shacl_from_file(client, shacl_rules_file_external_id)
            print(f"  Loaded {len(shacl_rules)} characters from file")
        except Exception as e:
            return {
                "conforms": False,
                "error": f"Failed to load SHACL rules from file '{shacl_rules_file_external_id}': {e!s}",
                "violations_found": 0,
            }

    if not shacl_rules:
        return {
            "conforms": False,
            "error": "Missing required parameter: shacl_rules or shacl_rules_file_external_id",
            "violations_found": 0,
        }

    # Parse SHACL rules graph once for blank-node sourceShape resolution
    from rdflib import Graph as _RawGraph

    _shacl_rules_graph = _RawGraph()
    try:
        _shacl_rules_graph.parse(data=shacl_rules, format="turtle")
    except Exception:
        pass  # best-effort; resolution will silently skip if graph is empty

    # 1. Load state from RawValidationState container
    print(f"Loading validation state for {db_name}.{table_name}...")
    state = _load_raw_validation_state(client, db_name, table_name, data_quality_space)
    min_timestamp = state["last_processed_timestamp"] if state else None

    if min_timestamp:
        print(f"  Last processed timestamp: {min_timestamp}")
    else:
        print("  No previous state found - validating all rows")

    # 2. Validate rows with SHACL
    print("Validating RAW table rows...")
    try:
        neat = NeatSession(client)
        result = neat.validate_instances.with_shacl_raw(
            db_name=db_name,
            table_name=table_name,
            shacl_rules=shacl_rules,
            min_last_updated_time=min_timestamp,
            limit=batch_size,
            verbose=verbose,
        )

        print(f"  Validation {'PASSED' if result.conforms else 'FAILED'}")

        # 3. Parse violations and post to Records API
        violations = _parse_shacl_report(result.report_graph, shacl_graph=_shacl_rules_graph)
        print(f"  Found {len(violations)} violations")

        # Post to Records API (if configured)
        stream_id = records_config.get("stream_id")
        rule_set_id = records_config.get("rule_set_id")
        rule_set_version = records_config.get("rule_set_version")

        records_posted = 0
        all_rows = []
        if stream_id and rule_set_id and rule_set_version:
            print("=" * 80)
            print("POSTING TO RECORDS API")
            print("=" * 80)

            # Build all_instances list (all rows that were validated)
            # For RAW, we need to fetch the rows again to get full row data
            from cognite_data_quality._neat._constants import get_raw_namespace

            all_rows = list(client.raw.rows.list(db_name, table_name, limit=batch_size))
            # For RAW validation, space is empty since RAW rows are not DMS instances
            # The focus node URI will be: {namespace_base}{row_key} (no space in between)
            all_instances = [
                {
                    "externalId": row.key,
                    "space": "",  # Empty for RAW (not a DMS space)
                    **row.columns,
                }
                for row in all_rows
            ]

            namespace_base = get_raw_namespace(db_name, table_name)

            posted_count, record_errors = post_validation_results_to_records(
                client=client,
                all_instances=all_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_config.get("records_space", "dataQuality"),
                records_container=records_config.get("records_container", "DataQualityValidationRecord"),
            )
            records_posted = posted_count
            print(f"  Posted {posted_count} records")
            if record_errors:
                print(f"  WARNING: {len(record_errors)} record(s) failed to post")
        else:
            print("  Skipping Records API (stream_id, rule_set_id, or rule_set_version not configured)")

        # 4. Update state with new timestamp (current time)
        new_timestamp = int(time.time() * 1000)  # Current time in milliseconds
        # Count actual rows validated (not violations)
        rows_count = len(all_rows)
        violation_count = len(violations)

        _save_raw_validation_state(
            client=client,
            db_name=db_name,
            table_name=table_name,
            last_processed_timestamp=new_timestamp,
            rows_processed=rows_count,
            conforms=result.conforms,
            violation_count=violation_count,
            data_quality_space=data_quality_space,
        )
        print(f"  Updated state with new timestamp: {new_timestamp}")

        return {
            "conforms": result.conforms,
            "rows_validated": rows_count,
            "violations_found": violation_count,
            "last_processed_timestamp": new_timestamp,
            "records_posted": records_posted,
        }

    except Exception as e:
        print(f"ERROR during validation: {e!s}")
        import traceback

        traceback.print_exc()
        return {
            "conforms": False,
            "error": str(e),
            "violations_found": 0,
        }


def handle_raw_validation_historic(data: dict[str, Any], client) -> dict[str, Any]:
    """
    Validate RAW table rows using cursor-based partitioning (historic/backfill).

    Used for initial validation of existing data or backfill operations.
    Processes a partition of historic data using a cursor from getCursors API.

    Implements 480-second timeout with cursor state saving for resumption.

    Args:
        data: Input data containing:
            - db_name: RAW database name
            - table_name: RAW table name
            - partition_id: Partition identifier for state tracking
            - cursor: Cursor string from getCursors API (or resume from saved state)
            - orchestration_id: (optional) Orchestration ID for state tracking
            - shacl_rules: SHACL rules as Turtle string (use this OR shacl_rules_file_external_id)
            - shacl_rules_file_external_id: External ID of a CDF File containing SHACL rules
            - verbose: (optional) Print progress messages (default: True)
            - job_run_id: (optional) Unique job run ID for records
            - records_config: Config for Records API posting
        client: Cognite client instance

    Returns:
        Dict with validation results:
            - conforms: Boolean indicating if validation passed
            - rows_validated: Number of rows validated
            - violations_found: Number of violations found
            - records_posted: Number of records posted to Records API
            - completed: Boolean indicating if partition is fully processed
            - next_cursor: Cursor for resumption if not completed
    """
    print("=" * 80)
    print("RAW TABLE SHACL VALIDATION (HISTORIC/CURSOR)")
    print("=" * 80)

    # Extract parameters
    db_name = data.get("db_name")
    table_name = data.get("table_name")
    partition_id = data.get("partition_id", "default")
    orchestration_id = data.get("orchestration_id")
    initial_cursor = data.get("cursor")
    shacl_rules = data.get("shacl_rules")
    shacl_rules_file_external_id = data.get("shacl_rules_file_external_id")
    verbose = data.get("verbose", True)
    data_quality_space = data.get("data_quality_space", "dataQuality")

    # Records API parameters
    job_run_id = data.get("job_run_id") or orchestration_id or f"job_{int(time.time())}"
    records_config = data.get("records_config", {})

    # Timeout configuration (480 seconds = 8 minutes, same as DMS)
    TIMEOUT_SECONDS = 480
    CURSOR_SAVE_INTERVAL = 30  # Save cursor state every 30 seconds

    # Validate required parameters
    if not db_name or not table_name:
        return {
            "conforms": False,
            "error": "Missing required parameters: db_name, table_name",
            "violations_found": 0,
            "completed": False,
        }

    # Try to resume from saved cursor state if no initial cursor provided
    if not initial_cursor and partition_id:
        print(f"Attempting to resume from saved state for partition: {partition_id}")
        saved_state = _load_partition_cursor_state(client, db_name, table_name, partition_id, data_quality_space)
        if saved_state:
            initial_cursor = saved_state.get("cursor")
            print(f"  Resuming from saved cursor (processed: {saved_state.get('processed_count', 0)} rows)")

    if not initial_cursor:
        return {
            "conforms": False,
            "error": "Missing required parameter: cursor (and no saved state found)",
            "violations_found": 0,
            "completed": False,
        }

    # Load SHACL rules from CDF File if external ID is provided
    if shacl_rules_file_external_id and not shacl_rules:
        try:
            print(f"Loading SHACL rules from CDF File: {shacl_rules_file_external_id}")
            shacl_rules = load_shacl_from_file(client, shacl_rules_file_external_id)
            print(f"  Loaded {len(shacl_rules)} characters from file")
        except Exception as e:
            return {
                "conforms": False,
                "error": f"Failed to load SHACL rules from file '{shacl_rules_file_external_id}': {e!s}",
                "violations_found": 0,
                "completed": False,
            }

    if not shacl_rules:
        return {
            "conforms": False,
            "error": "Missing required parameter: shacl_rules or shacl_rules_file_external_id",
            "violations_found": 0,
            "completed": False,
        }

    # Parse SHACL rules graph once for blank-node sourceShape resolution
    from rdflib import Graph as _RawGraph

    _shacl_rules_graph = _RawGraph()
    try:
        _shacl_rules_graph.parse(data=shacl_rules, format="turtle")
    except Exception:
        pass  # best-effort; resolution will silently skip if graph is empty

    # Process with timeout and cursor state management
    print(f"Validating RAW table partition (timeout: {TIMEOUT_SECONDS}s)...")
    print(f"  Partition ID: {partition_id}")
    print(
        f"  Initial cursor: {initial_cursor[:50]}..."
        if len(initial_cursor) > 50
        else f"  Initial cursor: {initial_cursor}"
    )

    start_time = time_module.time()
    last_cursor_save = time_module.time()
    total_rows = 0
    total_violations = 0
    current_cursor = initial_cursor
    completed = False

    try:
        # Process rows in batches with timeout management
        while True:
            # Check timeout
            elapsed = time_module.time() - start_time
            if elapsed >= TIMEOUT_SECONDS:
                print(f"  Timeout reached ({elapsed:.1f}s), saving state for resumption...")
                _save_partition_cursor_state(
                    client=client,
                    db_name=db_name,
                    table_name=table_name,
                    partition_id=partition_id,
                    cursor=current_cursor,
                    processed_count=total_rows,
                    orchestration_id=orchestration_id,
                    data_quality_space=data_quality_space,
                )
                return {
                    "conforms": total_violations == 0,
                    "rows_validated": total_rows,
                    "violations_found": total_violations,
                    "records_posted": total_rows,  # Approximate
                    "completed": False,
                    "next_cursor": current_cursor,
                    "elapsed_seconds": elapsed,
                }

            # Validate batch with current cursor
            neat = NeatSession(client)
            result = neat.validate_instances.with_shacl_raw(
                db_name=db_name,
                table_name=table_name,
                shacl_rules=shacl_rules,
                cursor=current_cursor,
                limit=1000,  # Process in smaller batches for timeout management
                verbose=False,  # Reduce verbosity for batch processing
            )

            # Parse violations and post to Records API
            violations = _parse_shacl_report(result.report_graph, shacl_graph=_shacl_rules_graph)
            batch_size = len(violations)
            batch_violations = len([v for v in violations if not v.get("conforms", True)])

            if batch_size > 0:
                # Post to Records API (if configured)
                stream_id = records_config.get("stream_id")
                rule_set_id = records_config.get("rule_set_id")
                rule_set_version = records_config.get("rule_set_version")

                if stream_id and rule_set_id and rule_set_version:
                    # Build all_instances list from validated batch
                    # Fetch the same batch of rows to get actual data
                    from cognite_data_quality._neat._constants import get_raw_namespace
                    from cognite_data_quality._neat._extractors._raw import RAWExtractor

                    extractor = RAWExtractor(
                        client=client,
                        db_name=db_name,
                        table_name=table_name,
                        cursor=current_cursor,
                        limit=1000,
                    )

                    # Extract rows (not triples)
                    all_rows = []
                    for row in extractor._fetch_filtered_rows():
                        all_rows.append(row)

                    # For RAW validation, space is empty since RAW rows are not DMS instances
                    all_instances = [
                        {
                            "externalId": row.key,
                            "space": "",  # Empty for RAW (not a DMS space)
                            **row.columns,
                        }
                        for row in all_rows
                    ]

                    namespace_base = get_raw_namespace(db_name, table_name)

                    _posted_count, _record_errors = post_validation_results_to_records(
                        client=client,
                        all_instances=all_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_config.get("records_space", "dataQuality"),
                        records_container=records_config.get("records_container", "DataQualityValidationRecord"),
                    )

                total_rows += batch_size
                total_violations += batch_violations

                if verbose or total_rows % 5000 == 0:
                    print(f"  Processed: {total_rows} rows, {total_violations} violations")

            # Check if we need to save cursor state
            if time_module.time() - last_cursor_save >= CURSOR_SAVE_INTERVAL:
                _save_partition_cursor_state(
                    client=client,
                    db_name=db_name,
                    table_name=table_name,
                    partition_id=partition_id,
                    cursor=current_cursor,
                    processed_count=total_rows,
                    orchestration_id=orchestration_id,
                    data_quality_space=data_quality_space,
                )
                last_cursor_save = time_module.time()

            # Check if we're done (no more rows)
            if batch_size == 0:
                completed = True
                print(f"  Partition completed: {total_rows} rows processed")

                # Clean up saved state
                _delete_partition_cursor_state(client, db_name, table_name, partition_id, data_quality_space)
                break

            # Note: In a real pagination scenario, we'd get nextCursor from the API response
            # For now, we assume the cursor mode processes until completion
            # This would need to be updated based on actual API behavior
            break  # Remove this once proper pagination with nextCursor is implemented

        return {
            "conforms": total_violations == 0,
            "rows_validated": total_rows,
            "violations_found": total_violations,
            "records_posted": total_rows,
            "completed": completed,
            "elapsed_seconds": time_module.time() - start_time,
        }

    except Exception as e:
        print(f"ERROR during validation: {e!s}")
        import traceback

        traceback.print_exc()

        # Save state before failing
        _save_partition_cursor_state(
            client=client,
            db_name=db_name,
            table_name=table_name,
            partition_id=partition_id,
            cursor=current_cursor,
            processed_count=total_rows,
            orchestration_id=orchestration_id,
            data_quality_space=data_quality_space,
        )

        return {
            "conforms": False,
            "error": str(e),
            "violations_found": total_violations,
            "completed": False,
            "next_cursor": current_cursor,
        }


# Helper functions


def _load_raw_validation_state(
    client, db_name: str, table_name: str, data_quality_space: str = "dataQuality"
) -> dict | None:
    """Load state for a RAW table, returns None if first run."""
    try:
        from cognite.client import data_modeling as dm

        # Query RawValidationState container
        external_id = f"{db_name}_{table_name}"
        result = client.data_modeling.instances.retrieve(
            nodes=[(dm.NodeId(data_quality_space, external_id))],
        )

        if result.nodes:
            node = result.nodes[0]
            props = node.properties.get(data_quality_space, {}).get("RawValidationState/1", {})
            return {
                "last_processed_timestamp": props.get("last_processed_timestamp"),
                "last_run_time": props.get("last_run_time"),
                "rows_processed": props.get("rows_processed", 0),
                "last_validation_conforms": props.get("last_validation_conforms") == "true",
                "last_violation_count": props.get("last_violation_count", 0),
            }
    except Exception as e:
        print(f"  Warning: Could not load state: {e}")

    return None


def _save_raw_validation_state(
    client,
    db_name: str,
    table_name: str,
    last_processed_timestamp: int,
    rows_processed: int,
    conforms: bool,
    violation_count: int,
    data_quality_space: str = "dataQuality",
):
    """Update state after successful validation."""
    from cognite.client import data_modeling as dm

    external_id = f"{db_name}_{table_name}"

    # Create or update instance in RawValidationState container
    from datetime import datetime, timezone

    # Format current time as ISO 8601 string for Timestamp type
    current_time = datetime.now(timezone.utc).isoformat(timespec="milliseconds")

    instance = dm.NodeApply(
        space=data_quality_space,
        external_id=external_id,
        sources=[
            dm.NodeOrEdgeData(
                source=dm.ContainerId(data_quality_space, "RawValidationState"),
                properties={
                    "db_name": db_name,
                    "table_name": table_name,
                    "last_processed_timestamp": last_processed_timestamp,
                    "last_run_time": current_time,
                    "rows_processed": rows_processed,
                    "last_validation_conforms": "true" if conforms else "false",
                    "last_violation_count": violation_count,
                },
            )
        ],
    )

    client.data_modeling.instances.apply(instance)


def _parse_shacl_report(report_graph, shacl_graph=None) -> list[dict]:
    """
    Parse SHACL validation report graph to extract violations.

    Args:
        report_graph: Either an rdflib.Graph object or a Turtle string
        shacl_graph: Parsed SHACL rules graph used to resolve blank-node sourceShape
                     values to their parent named NodeShape URI.

    Returns list of violation dicts with instance information.
    """
    from rdflib import BNode, Graph
    from rdflib.namespace import SH

    # Build path → parent named shape mapping from SHACL rules graph
    _path_to_shape: dict = {}
    if shacl_graph is not None:
        try:
            for _parent, _prop_bnode in shacl_graph.subject_objects(predicate=SH.property):
                if not isinstance(_parent, BNode):
                    _prop_path = shacl_graph.value(_prop_bnode, SH.path)
                    if _prop_path and not isinstance(_prop_path, BNode):
                        _path_to_shape[str(_prop_path)] = str(_parent)
        except Exception:
            pass

    violations = []

    try:
        # Handle both Graph objects and strings
        if isinstance(report_graph, Graph):
            g = report_graph
        else:
            g = Graph()
            g.parse(data=report_graph, format="turtle")

        # Find all validation results
        for result in g.subjects(predicate=SH.resultSeverity):
            focus_node = g.value(result, SH.focusNode)
            result_message = g.value(result, SH.resultMessage)
            source_constraint = g.value(result, SH.sourceConstraintComponent)
            result_path = g.value(result, SH.resultPath)
            value = g.value(result, SH.value)
            source_shape = g.value(result, SH.sourceShape)
            severity = g.value(result, SH.resultSeverity)

            # Blank node sourceShape can't be matched across graph re-parsings;
            # resolve via resultPath → parent named shape instead
            if isinstance(source_shape, BNode):
                source_shape = None
            if source_shape is None and result_path and str(result_path) in _path_to_shape:
                source_shape = _path_to_shape[str(result_path)]

            # Keep the full focus node URI for violation matching
            focus_uri = str(focus_node) if focus_node else ""

            violations.append(
                {
                    "conforms": False,
                    "focusNode": focus_uri,  # Full URI for matching with instances_to_post
                    "resultMessage": str(result_message) if result_message else "Validation failed",
                    "sourceConstraintComponent": str(source_constraint) if source_constraint else None,
                    "resultPath": str(result_path) if result_path else None,
                    "value": str(value) if value else None,
                    "sourceShape": str(source_shape) if source_shape else None,
                    "resultSeverity": str(severity) if severity else "Violation",
                }
            )
    except Exception as e:
        print(f"  Warning: Could not parse SHACL report: {e}")

    return violations


def _load_partition_cursor_state(
    client, db_name: str, table_name: str, partition_id: str, data_quality_space: str = "dataQuality"
) -> dict | None:
    """Load saved cursor state for a partition."""
    try:
        from cognite.client import data_modeling as dm

        # Use FunctionValidationState container (same as DMS partitioned validation)
        external_id = f"raw_{db_name}_{table_name}_{partition_id}"
        result = client.data_modeling.instances.retrieve(
            nodes=[(dm.NodeId(data_quality_space, external_id))],
        )

        if result.nodes:
            node = result.nodes[0]
            props = node.properties.get(data_quality_space, {}).get("FunctionValidationState/1", {})
            return {
                "cursor": props.get("syncCursor"),
                "processed_count": props.get("processedCount", 0),
                "last_updated": props.get("lastUpdated"),
            }
    except Exception as e:
        print(f"  Warning: Could not load cursor state: {e}")

    return None


def _save_partition_cursor_state(
    client,
    db_name: str,
    table_name: str,
    partition_id: str,
    cursor: str,
    processed_count: int,
    orchestration_id: str | None = None,
    data_quality_space: str = "dataQuality",
):
    """Save cursor state for resumption after timeout."""
    from cognite.client import data_modeling as dm

    external_id = f"raw_{db_name}_{table_name}_{partition_id}"

    # Create or update instance in FunctionValidationState container
    instance = dm.NodeApply(
        space=data_quality_space,
        external_id=external_id,
        sources=[
            dm.NodeOrEdgeData(
                source=dm.ContainerId(data_quality_space, "FunctionValidationState"),
                properties={
                    "orchestration_id": orchestration_id or f"raw_{db_name}_{table_name}",
                    "partitionId": partition_id,
                    "syncCursor": cursor,
                    "processedCount": processed_count,
                    "lastUpdated": int(time.time() * 1000),
                },
            )
        ],
    )

    client.data_modeling.instances.apply(instance)
    print(f"  Saved cursor state: {processed_count} rows processed")


def _delete_partition_cursor_state(
    client, db_name: str, table_name: str, partition_id: str, data_quality_space: str = "dataQuality"
):
    """Delete cursor state after successful completion."""
    try:
        from cognite.client import data_modeling as dm

        external_id = f"raw_{db_name}_{table_name}_{partition_id}"
        client.data_modeling.instances.delete(
            nodes=[(dm.NodeId(data_quality_space, external_id))],
        )
        print(f"  Cleaned up cursor state for partition: {partition_id}")
    except Exception as e:
        print(f"  Warning: Could not delete cursor state: {e}")
