"""
Cognite Function handler for time series SHACL validation.

Validates time series data quality using SHACL rules with CDF SPARQL functions.

Modes:
1. Normal mode: SHACL rules use their specified time windows (e.g., 90m-ago)
2. Backfill mode: Process historic data in chunks matching the window size

Time windows should be configured in SHACL rules. For a 60-minute schedule,
use a 90-minute window for 30-minute overlap to catch edge cases.

Results are posted to CDF Records API with deduplication based on violation hash.
"""

from __future__ import annotations

import datetime
import logging
import re
from typing import Any

from cognite.client import CogniteClient
from cognite.client.data_classes.data_modeling import NodeId
from cognite.client.data_classes.filters import And, Equals, HasData, Prefix
from common.cursor_state import (
    ensure_cursor_container,
    load_cursor_state,
    save_cursor_state,
)
from common.time_window import get_aligned_hour_window, replace_time_windows_in_shacl

logger = logging.getLogger(__name__)

_VIOLATION_SEVERITY = "violation"


def _inject_prefetch_errors(
    instances: list[dict[str, Any]],
    violations_by_ts: dict[str, list[dict[str, Any]]],
) -> None:
    """Inject DataFetchError violations for time series where pre-fetch permanently failed."""
    from cognite_data_quality._neat._sparql_functions._sdk_wrappers import get_prefetch_errors

    fetch_errors = get_prefetch_errors()
    for inst in instances:
        ts_id = f"{inst['space']}/{inst['externalId']}"
        if ts_id in fetch_errors:
            violations_by_ts.setdefault(ts_id, []).append(
                {
                    "sourceConstraintComponent": "DataFetchError",
                    "severity": "Warning",
                    "sourceShape": "PrefetchError",
                    "message": (
                        "Validation infrastructure error: failed to fetch time series data "
                        "from CDF after retries. Validation result may be incomplete."
                    ),
                }
            )


# ============================================================================
# Time Series Fetching
# ============================================================================


def _fetch_timeseries_by_filter(
    client: CogniteClient,
    filter_config: dict[str, Any],
    space: str = "cdf_cdm",
    view_external_id: str = "CogniteTimeSeries",
    view_version: str = "v1",
) -> list[dict[str, Any]]:
    """Fetch CogniteTimeSeries instances using DMS API filter."""
    filters = []
    view_id = (space, view_external_id, view_version)
    filters.append(HasData(views=[view_id]))

    if "space" in filter_config:
        filters.append(Equals(["node", "space"], filter_config["space"]))
    if "external_id_prefix" in filter_config:
        filters.append(Prefix(["node", "externalId"], filter_config["external_id_prefix"]))

    combined_filter = And(*filters) if len(filters) > 1 else filters[0]

    nodes = client.data_modeling.instances.list(
        instance_type="node",
        filter=combined_filter,
        limit=1000,
    )

    return [{"externalId": node.external_id, "space": node.space, "properties": {}} for node in nodes]


def _fetch_timeseries_by_ids(
    client: CogniteClient,
    instance_ids: list[dict[str, str]],
) -> list[dict[str, Any]]:
    """Fetch CogniteTimeSeries instances by explicit instance IDs."""
    node_ids = [NodeId(space=item["space"], external_id=item["external_id"]) for item in instance_ids]

    result = client.data_modeling.instances.retrieve(nodes=node_ids)

    return [{"externalId": node.external_id, "space": node.space, "properties": {}} for node in result.nodes]


# ============================================================================
# DMS Container Setup
# ============================================================================

RESULTS_CONTAINER_ID = "TimeSeriesValidationResults"


def _ensure_results_container(client: CogniteClient, records_config: dict[str, Any]) -> None:
    """Create the results storage container if it doesn't exist (used by Records API)."""
    from cognite.client.data_classes.data_modeling import (
        ContainerApply,
        ContainerProperty,
        SpaceApply,
        Text,
    )

    space = records_config.get("records_space", "dataQuality")
    container_id = records_config.get("records_container", RESULTS_CONTAINER_ID)

    try:
        client.data_modeling.spaces.apply(SpaceApply(space=space))
    except Exception:
        pass

    container = ContainerApply(
        space=space,
        external_id=container_id,
        name="Time Series Validation Results",
        description="Stores validation results for time series (Records API)",
        properties={
            "timeseriesId": ContainerProperty(type=Text(), nullable=False),
            "validatedAt": ContainerProperty(type=Text(), nullable=False),
            "conforms": ContainerProperty(type=Text(), nullable=False),
            "status": ContainerProperty(type=Text(), nullable=False),
            "violationCount": ContainerProperty(type=Text(), nullable=False),
            "violations": ContainerProperty(type=Text(), nullable=True),
            "timeWindow": ContainerProperty(type=Text(), nullable=True),
            "ruleSetId": ContainerProperty(type=Text(), nullable=False),
        },
    )

    try:
        client.data_modeling.containers.apply(container)
        logger.info(f"Results container ready: {space}/{container_id}")
    except Exception as e:
        logger.debug(f"Results container might already exist: {e}")


# ============================================================================
# Validation Result Processing
# ============================================================================


def _extract_violations_from_report(report_text: str) -> list[dict[str, Any]]:
    """Extract violation details from SHACL report text.

    pyshacl report format example:
        Constraint Violation in MinCountConstraintComponent (http://www.w3.org/ns/shacl#MinCountConstraintComponent):
            Severity: sh:Violation
            Source Shape: ex:DatapointsCountShape
            Focus Node: <https://cdf.cognite.com/ts_dq_test/dq_test_empty_1>
            Result Path: cdf_sdk:datapoints_count
            Message: Time series must have at least 1 datapoint
            Value: 0
    """
    violations = []
    current_violation: dict[str, str] = {}

    for line in report_text.split("\n"):
        line = line.strip()
        # New violation starts with "Constraint Violation" or "Validation Result"
        if line.startswith("Constraint Violation") or line.startswith("Validation Result"):
            if current_violation:
                violations.append(current_violation)
            current_violation = {}
            # Extract constraint component from the header line
            if "(" in line and ")" in line:
                constraint = line.split("(")[-1].rstrip("):")
                current_violation["sourceConstraintComponent"] = constraint
        elif ": " in line:
            key, value = line.split(": ", 1)
            key = key.strip()
            value = value.strip()
            # Normalize key names
            if key == "Severity":
                # sh:Violation -> Violation
                current_violation["severity"] = value.split(":")[-1] if ":" in value else value
            elif key == "Source Shape":
                current_violation["sourceShape"] = value
            elif key == "Focus Node":
                current_violation["focusNode"] = value.strip("<>")
            elif key == "Result Path":
                current_violation["resultPath"] = value
            elif key == "Message":
                current_violation["message"] = value
            elif key == "Value":
                current_violation["value"] = value
            else:
                current_violation[key] = value

    if current_violation:
        violations.append(current_violation)

    return violations


def _post_validation_result(
    client: CogniteClient,
    timeseries_id: str,
    timeseries_space: str,
    conforms: bool,
    violations: list[dict[str, Any]],
    records_config: dict[str, Any],
    time_window: str | None = None,
) -> bool:
    """Post validation result to Records API using same schema as instance validation."""
    datetime.datetime.now(tz=datetime.timezone.utc).isoformat()
    space = records_config.get("records_space", "dataQuality")
    container = records_config.get("records_container", "DataQualityValidationRecord")
    stream_id = records_config.get("stream_id", "dq_validation_stream")
    rule_set_id = records_config.get("rule_set_id", "TimeSeriesSHACL")
    rule_set_version = records_config.get("rule_set_version", "1.0")
    job_run_id = records_config.get("job_run_id") or f"job_{int(datetime.datetime.now().timestamp())}"

    # Build constraint details from violations
    failed_constraints = []
    severities = []
    constraint_details = []
    source_shapes = set()

    for v in violations:
        constraint = v.get("sourceConstraintComponent", "Unknown")
        failed_constraints.append(constraint)
        severity = v.get("severity", "Violation")
        if severity not in severities:
            severities.append(severity)
        if v.get("sourceShape"):
            source_shapes.add(v.get("sourceShape"))
        constraint_details.append(
            {
                "constraint": constraint,
                "message": v.get("message", ""),
                "severity": severity,
                "value": v.get("value", ""),
                "resultPath": v.get("resultPath", ""),
                "sourceShape": v.get("sourceShape", ""),
            }
        )

    # Focus node in same format as instance validation
    focus_node = f"https://cdf.cognite.com/{timeseries_space}/{timeseries_id}"
    record_external_id = f"{timeseries_space}/{timeseries_id}/{job_run_id}"
    if time_window:
        record_external_id = f"{timeseries_space}/{timeseries_id}/{time_window}"

    record = {
        "items": [
            {
                "space": space,
                "externalId": record_external_id,
                "sources": [
                    {
                        "source": {"type": "container", "space": space, "externalId": container},
                        "properties": {
                            "ruleSetId": rule_set_id,
                            "ruleSetVersion": rule_set_version,
                            "jobRunId": job_run_id,
                            "passedValidation": not any(s.lower() == _VIOLATION_SEVERITY for s in severities),
                            "resultSeverity": severities if severities else [],
                            "failedConstraints": failed_constraints,
                            "focusNode": focus_node,
                            "focusNodeInstance": {"space": timeseries_space, "externalId": timeseries_id},
                            "validationReport": {
                                "violationCount": len(violations),
                                "violations": constraint_details,
                                "summary": f"{len(violations)} violation(s)"
                                if violations
                                else "Passed all constraints",
                            },
                        },
                    }
                ],
            }
        ]
    }

    if source_shapes:
        record["items"][0]["sources"][0]["properties"]["sourceShape"] = list(source_shapes)

    # Enable alpha version for Records API
    original_headers = client.config.headers.copy()
    client.config.headers["cdf-version"] = "alpha"

    try:
        client.post(f"/api/v1/projects/{client.config.project}/streams/{stream_id}/records", json=record)
        status = "PASSED" if conforms else "FAILED"
        print(f"  Posted [{status}] record for {timeseries_id}")
        return True
    except Exception as e:
        print(f"  ERROR posting record for {timeseries_space}/{timeseries_id}: {e}")
        return False
    finally:
        client.config.headers = original_headers


# ============================================================================
# Validation Runner
# ============================================================================


def _run_validation(
    client: CogniteClient,
    instances: list[dict[str, Any]],
    shacl_rules: str,
    datamodel_space: str,
    datamodel_external_id: str,
    datamodel_version: str,
    verbose: bool = True,
    backfill_start: datetime.datetime | None = None,
    backfill_end: datetime.datetime | None = None,
) -> tuple[bool, str, str]:
    """Run SHACL validation.

    Args:
        backfill_start: Optional start time for INDSL data fetching (ms timestamp or ISO string)
        backfill_end: Optional end time for INDSL data fetching (ms timestamp or ISO string)
    """
    from cognite_data_quality._neat import NeatSession

    neat = NeatSession(client)

    # Convert datetime to ISO string for NEAT
    bf_start = backfill_start.isoformat() if backfill_start else None
    bf_end = backfill_end.isoformat() if backfill_end else None

    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,
        default_view_space=datamodel_space,
        default_view_name=datamodel_external_id,
        auto_load_depth=0,
        verbose=verbose,
        backfill_start=bf_start,
        backfill_end=bf_end,
    )

    return conforms, report_graph, report_text


# ============================================================================
# Main Handler
# ============================================================================


def handle_timeseries_validation(data: dict[str, Any], client: CogniteClient) -> dict[str, Any]:
    """
    Main handler for time series SHACL validation.

    Supports two modes:
    1. Normal mode: Run validation with SHACL-specified time windows
    2. Backfill mode: Process historic data in chunks

    Args:
        data: Configuration dict containing:
            - shacl_rules_file_external_id: File external_id for SHACL rules
            - datamodel_space/external_id/version: DMS view info
            - filter or instance_ids: Time series selection
            - backfill: Optional backfill config
              - enabled: true to enable backfill
              - start_time: ISO timestamp or "30d-ago"
              - end_time: ISO timestamp or "now"
              - window_minutes: Chunk size (should match SHACL rule window)
            - records_config: Records API config
            - verbose: Print progress (default: True)
        client: CogniteClient

    Returns:
        Dict with validation results summary
    """
    verbose = data.get("verbose", True)
    data_quality_space = data.get("data_quality_space", "dataQuality")

    if verbose:
        print("=" * 60)
        print("TIME SERIES SHACL VALIDATION")
        print("=" * 60)

    # Extract config
    shacl_rules_file_id = data.get("shacl_rules_file_external_id")
    datamodel_space = data.get("datamodel_space", "cdf_cdm")
    datamodel_external_id = data.get("datamodel_external_id", "CogniteTimeSeries")
    datamodel_version = data.get("datamodel_version", "v1")
    records_config = data.get("records_config", {})
    if "job_run_id" not in records_config and data.get("job_run_id"):
        records_config = {**records_config, "job_run_id": data["job_run_id"]}
    backfill_config = data.get("backfill", {})

    is_backfill = backfill_config.get("enabled", False)

    if verbose:
        print(f"Mode: {'BACKFILL' if is_backfill else 'NORMAL'}")

    # 1. Fetch time series instances
    if verbose:
        print("[1/4] Fetching time series instances...")

    if "instance_ids" in data:
        instances = _fetch_timeseries_by_ids(client, data["instance_ids"])
    elif "filter" in data:
        instances = _fetch_timeseries_by_filter(
            client,
            data["filter"],
            datamodel_space,
            datamodel_external_id,
            datamodel_version,
        )
    else:
        raise ValueError("Either 'instance_ids' or 'filter' must be provided")

    if verbose:
        print(f"  Found {len(instances)} time series")

    if not instances:
        return {"status": "success", "message": "No time series found", "timeseries_count": 0}

    # 2. Load SHACL rules
    if verbose:
        print("[2/4] Loading SHACL rules...")

    shacl_file = client.files.retrieve(external_id=shacl_rules_file_id)
    if not shacl_file:
        raise ValueError(f"SHACL rules file not found: {shacl_rules_file_id}")

    shacl_rules_original = client.files.download_bytes(external_id=shacl_rules_file_id).decode("utf-8")

    if verbose:
        print(f"  Loaded rules from: {shacl_rules_file_id}")

    # 3. Setup storage containers
    if verbose:
        print("[3/4] Setting up storage...")

    ensure_cursor_container(client, data_quality_space)
    _ensure_results_container(client, records_config)
    ts_ids = [f"{inst['space']}/{inst['externalId']}" for inst in instances]
    rule_set_id = records_config.get("rule_set_id", "TimeSeriesSHACL")
    cursor_state = load_cursor_state(client, ts_ids, rule_set_id, data_quality_space)

    if verbose:
        print(f"  Loaded state for {len(cursor_state)} time series")

    # 4. Run validation
    results: dict[str, Any] = {
        "status": "success",
        "mode": "backfill" if is_backfill else "normal",
        "timeseries_count": len(instances),
        "chunks_processed": 0,
        "total_violations": 0,
        "records_posted": 0,
        "records_skipped": 0,
    }

    if is_backfill:
        # Backfill mode: process in chunks aligned to hour boundaries
        if verbose:
            print("[4/4] Running backfill validation...")

        window_minutes = backfill_config.get("window_minutes", 60)
        start_str = backfill_config.get("start_time", "30d-ago")
        end_str = backfill_config.get("end_time", "now")

        # Parse start/end times
        now = datetime.datetime.now(tz=datetime.timezone.utc)

        if start_str.endswith("-ago"):
            match = re.match(r"(\d+)([mhd])-ago", start_str)
            if match:
                value, unit = int(match.group(1)), match.group(2)
                delta = {
                    "m": datetime.timedelta(minutes=value),
                    "h": datetime.timedelta(hours=value),
                    "d": datetime.timedelta(days=value),
                }[unit]
                start_time = now - delta
            else:
                start_time = now - datetime.timedelta(days=30)
        else:
            start_time = datetime.datetime.fromisoformat(start_str)

        end_time = now if end_str == "now" else datetime.datetime.fromisoformat(end_str)

        # Align to hour boundaries
        start_time = start_time.replace(minute=0, second=0, microsecond=0)
        end_time = end_time.replace(minute=0, second=0, microsecond=0)

        # Check cursor state for resume point (use the most recent backfillProgress)
        resume_from = None
        for _ts_id, state in cursor_state.items():
            progress = state.get("backfillProgress")
            if progress:
                try:
                    progress_time = datetime.datetime.fromisoformat(progress)
                    if resume_from is None or progress_time > resume_from:
                        resume_from = progress_time
                except (ValueError, TypeError):
                    pass

        # Resume from last progress if available
        if resume_from and resume_from > start_time:
            if verbose:
                print(f"  Resuming from: {resume_from.strftime('%Y-%m-%d %H:%M')} UTC")
            start_time = resume_from

        # Check if backfill is already complete - switch to normal mode
        if start_time >= end_time:
            if verbose:
                print(
                    f"  Backfill complete (progress {start_time.strftime('%Y-%m-%d %H:%M')}"
                    f" >= end {end_time.strftime('%Y-%m-%d %H:%M')})"
                )
                print("  Switching to NORMAL mode...")
            is_backfill = False  # Fall through to normal mode below
        else:
            # Process in chunks
            chunk_delta = datetime.timedelta(minutes=window_minutes)
            current_start = start_time

            while current_start < end_time:
                current_end = min(current_start + chunk_delta, end_time)

                if verbose:
                    print(
                        f"\n  Chunk: {current_start.strftime('%Y-%m-%d %H:%M')}"
                        f" to {current_end.strftime('%Y-%m-%d %H:%M')} UTC"
                    )

                # Modify SHACL rules with actual timestamps
                shacl_rules = replace_time_windows_in_shacl(shacl_rules_original, current_start, current_end)

                # Run validation for this chunk
                _conforms, _, report_text = _run_validation(
                    client,
                    instances,
                    shacl_rules,
                    datamodel_space,
                    datamodel_external_id,
                    datamodel_version,
                    verbose=False,
                    backfill_start=current_start,
                    backfill_end=current_end,
                )

                # Process results
                violations = _extract_violations_from_report(report_text)
                results["chunks_processed"] += 1
                results["total_violations"] += len(violations)

                # Group violations by time series and post (INSIDE the loop)
                violations_by_ts: dict[str, list] = {}
                for v in violations:
                    focus = v.get("focusNode", "")
                    for inst in instances:
                        if inst["externalId"] in focus:
                            ts_id = f"{inst['space']}/{inst['externalId']}"
                            violations_by_ts.setdefault(ts_id, []).append(v)
                            break

                _inject_prefetch_errors(instances, violations_by_ts)

                for inst in instances:
                    ts_id = f"{inst['space']}/{inst['externalId']}"
                    ts_violations = violations_by_ts.get(ts_id, [])
                    fatal_violations = [v for v in ts_violations if v.get("severity", "Violation") == "Violation"]
                    ts_conforms = len(fatal_violations) == 0

                    posted = _post_validation_result(
                        client,
                        inst["externalId"],
                        inst["space"],
                        ts_conforms,
                        ts_violations,
                        records_config,
                        time_window=f"{current_start.isoformat()}/{current_end.isoformat()}",
                    )

                    if posted:
                        results["records_posted"] += 1
                    else:
                        results["records_skipped"] += 1

                # Save cursor state AFTER posting records (inside loop for continuous saving)
                # Save once per chunk with the latest timestamp
                cursor_save_errors = 0
                for inst in instances:
                    ts_id = f"{inst['space']}/{inst['externalId']}"
                    ts_violations = violations_by_ts.get(ts_id, [])
                    fatal_violations = [v for v in ts_violations if v.get("severity", "Violation") == "Violation"]
                    ts_conforms = len(fatal_violations) == 0
                    try:
                        save_cursor_state(
                            client,
                            ts_id,
                            records_config.get("rule_set_id", "TimeSeriesSHACL"),
                            "PASSED" if ts_conforms else "FAILED",
                            data_quality_space,
                            backfill_progress=current_end.isoformat(),
                        )
                    except Exception as e:
                        cursor_save_errors += 1
                        if cursor_save_errors == 1:  # Only log first error
                            print(f"    Cursor save error: {e}")

                if verbose:
                    print(
                        f"    Violations: {len(violations)}, Records posted: {results['records_posted']},"
                        f" Cursor saved to: {current_end.strftime('%Y-%m-%d %H:%M')}"
                    )

                current_start = current_end

    if not is_backfill:
        # Normal mode: validate the previous complete hour (aligned to hour boundaries)
        window_start, window_end = get_aligned_hour_window()
        time_window_str = f"{window_start.isoformat()}/{window_end.isoformat()}"

        if verbose:
            print("[4/4] Running validation...")
            print(f"  Time window: {window_start.strftime('%H:%M')} to {window_end.strftime('%H:%M')} UTC")

        # Replace time placeholders with actual timestamps
        shacl_rules = replace_time_windows_in_shacl(shacl_rules_original, window_start, window_end)

        _conforms, _, report_text = _run_validation(
            client,
            instances,
            shacl_rules,
            datamodel_space,
            datamodel_external_id,
            datamodel_version,
            verbose=verbose,
            backfill_start=window_start,
            backfill_end=window_end,
        )

        violations = _extract_violations_from_report(report_text)
        results["total_violations"] = len(violations)
        results["chunks_processed"] = 1

        # Group and post results
        violations_by_ts = {}
        for v in violations:
            focus = v.get("focusNode", "")  # Note: camelCase key from _extract_violations_from_report
            for inst in instances:
                if inst["externalId"] in focus:
                    ts_id = f"{inst['space']}/{inst['externalId']}"
                    violations_by_ts.setdefault(ts_id, []).append(v)
                    break

        _inject_prefetch_errors(instances, violations_by_ts)

        for inst in instances:
            ts_id = f"{inst['space']}/{inst['externalId']}"
            ts_violations = violations_by_ts.get(ts_id, [])
            fatal_violations = [v for v in ts_violations if v.get("severity", "Violation") == "Violation"]
            ts_conforms = len(fatal_violations) == 0

            posted = _post_validation_result(
                client,
                inst["externalId"],
                inst["space"],
                ts_conforms,
                ts_violations,
                records_config,
                time_window=time_window_str,
            )

            if posted:
                results["records_posted"] += 1
            else:
                results["records_skipped"] += 1

            save_cursor_state(
                client,
                ts_id,
                records_config.get("rule_set_id", "TimeSeriesSHACL"),
                "PASSED" if ts_conforms else "FAILED",
                data_quality_space,
            )

    if verbose:
        print("=" * 60)
        print("VALIDATION COMPLETE")
        print("=" * 60)
        print(f"  Mode: {results['mode']}")
        print(f"  Time series: {results['timeseries_count']}")
        print(f"  Chunks processed: {results['chunks_processed']}")
        print(f"  Total violations: {results['total_violations']}")
        print(f"  Records posted: {results['records_posted']}")
        print(f"  Records skipped: {results['records_skipped']}")

    return results
