"""
Handler for querying status of orchestration runs.

Returns a snapshot of active and completed historic validation runs by reading
from two DMS views:
  - OrchestrationStateView  — per-run summary (refreshed every ~2 min by monitor schedule)
  - FunctionValidationStateView — per-partition cursor progress for in-progress partitions

Input parameters:
  validation_type     : "orchestration_status"
  orchestration_id    : str  (optional) — scope to a single run
  include_stale       : bool (optional, default False) — include active runs not updated in >24h
  data_quality_space  : str  (optional, default "dataQuality")
"""

from __future__ import annotations

from datetime import datetime, timedelta, timezone
from typing import Any

from cognite.client import CogniteClient


def handle_orchestration_status(data: dict[str, Any], client: CogniteClient) -> dict[str, Any]:
    """
    Return status of orchestration runs.

    Categorises runs into:
      active_runs    — status == "monitoring", updated within 24h
      completed_runs — status == "completed" or "failed"
      stale_runs     — status == "monitoring", NOT updated within 24h
                       (only included when include_stale=True)
    """
    include_stale: bool = data.get("include_stale", False)
    filter_orch_id: str | None = data.get("orchestration_id")
    space: str = data.get("data_quality_space", "dataQuality")

    now_utc = datetime.now(timezone.utc)
    stale_cutoff = now_utc - timedelta(hours=24)

    try:
        orch_states = _list_orchestration_states(client, space, filter_orch_id)
    except Exception as e:
        return {"status": "error", "error": f"Failed to query OrchestrationStateView: {e}"}

    # Only fetch cursor states when there are active/stale monitoring runs
    active_orch_ids = [o["orchestration_id"] for o in orch_states if o.get("status") in ("monitoring", "in_progress")]
    cursor_states_by_orch: dict[str, dict[str, dict[str, Any]]] = {}
    if active_orch_ids:
        try:
            all_cursors = _list_partition_cursor_states(client, space)
            # Group by orchestration_id
            for partition_id, cursor_info in all_cursors.items():
                oid = cursor_info.get("orchestration_id", "")
                if oid not in cursor_states_by_orch:
                    cursor_states_by_orch[oid] = {}
                cursor_states_by_orch[oid][partition_id] = cursor_info
        except Exception as e:
            # Non-fatal: partition detail will be missing processedCount
            print(f"Warning: Failed to query FunctionValidationStateView: {e}")

    # Build summaries for all runs, then deduplicate: keep only the newest per view.
    # "newest" = latest last_updated timestamp.
    all_summaries: list[tuple[dict, str]] = []  # (summary, run_status)
    for orch in orch_states:
        oid = orch.get("orchestration_id", "")
        cursor_states = cursor_states_by_orch.get(oid, {})
        summary = _build_run_summary(orch, cursor_states, stale_cutoff)
        all_summaries.append((summary, orch.get("status", "")))

    # Deduplicate by view_external_id — keep newest last_updated per view.
    # Runs where the orchestrator found 0 instances have no status written — treat as completed.
    newest: dict[str, tuple[dict, str]] = {}
    for summary, run_status in all_summaries:
        if not run_status:
            run_status = "completed"
            summary["status"] = "completed"
        view_key = summary.get("view_external_id") or summary.get("orchestration_id", "")
        existing = newest.get(view_key)
        if existing is None or (summary.get("last_updated") or "") > (existing[0].get("last_updated") or ""):
            newest[view_key] = (summary, run_status)

    active_runs = []
    completed_runs = []
    stale_runs = []

    for summary, run_status in newest.values():
        if run_status in ("monitoring", "in_progress"):
            if summary["is_stale"]:
                stale_runs.append(summary)
            else:
                active_runs.append(summary)
        else:
            completed_runs.append(summary)

    # Sort: active/stale by last_updated descending, completed by last_updated descending
    active_runs.sort(key=lambda r: r.get("last_updated") or "", reverse=True)
    completed_runs.sort(key=lambda r: r.get("last_updated") or "", reverse=True)
    stale_runs.sort(key=lambda r: r.get("last_updated") or "", reverse=True)

    result: dict[str, Any] = {
        "status": "ok",
        "active_runs": active_runs,
        "completed_runs": completed_runs,
        "total_active": len(active_runs),
        "total_completed": len(completed_runs),
    }
    if include_stale:
        result["stale_runs"] = stale_runs
        result["total_stale"] = len(stale_runs)

    return result


def _list_orchestration_states(
    client: CogniteClient,
    data_quality_space: str,
    orchestration_id: str | None = None,
) -> list[dict[str, Any]]:
    """
    List all orchestration states from OrchestrationStateView.

    Optionally filter to a single orchestration_id (client-side).
    Pattern mirrors query_orchestration_states() in batch_dashboard.py.
    """
    from cognite.client.data_classes.data_modeling.ids import ViewId

    view_id = ViewId(data_quality_space, "OrchestrationStateView", "1")
    result = client.data_modeling.instances.list(
        instance_type="node",
        sources=view_id,
        limit=-1,
    )

    orchestrations = []
    for node in result:
        if not node.properties:
            continue

        props = node.properties.get(view_id)
        if not props:
            continue

        oid = props.get("orchestration_id", "")
        if orchestration_id and oid != orchestration_id:
            continue

        orchestrations.append(
            {
                "external_id": node.external_id,
                "orchestration_id": oid,
                "view_external_id": props.get("view_external_id", ""),
                "view_space": props.get("view_space", ""),
                "status": props.get("status", ""),
                "created_time": props.get("created_time"),
                "last_updated": props.get("last_updated"),
                "partition_count": props.get("partition_count", 0),
                "total_instances": props.get("total_instances", 0),
                "validated_instances": props.get("validated_instances", 0),
                "total_violations": props.get("total_violations", 0),
                # These are Json-typed DMS properties — SDK returns them as dicts
                "completed_partitions": props.get("completed_partitions") or {},
                "failed_partitions": props.get("failed_partitions") or {},
                "partition_call_ids": props.get("partition_call_ids") or {},
                "partition_retry_count": props.get("partition_retry_count") or {},
            }
        )

    return orchestrations


def _list_partition_cursor_states(
    client: CogniteClient,
    data_quality_space: str,
) -> dict[str, dict[str, Any]]:
    """
    List all partition cursor states from FunctionValidationStateView.

    Returns a flat dict keyed by partitionId with orchestration_id, processedCount,
    and lastUpdated for each active partition.

    Pattern mirrors query_cursor_states() in batch_dashboard.py.
    """
    from cognite.client.data_classes.data_modeling.ids import ViewId

    view_id = ViewId(data_quality_space, "FunctionValidationStateView", "1")
    result = client.data_modeling.instances.list(
        instance_type="node",
        sources=view_id,
        limit=-1,
    )

    cursors: dict[str, dict[str, Any]] = {}
    for node in result:
        if not node.properties:
            continue

        props = node.properties.get(view_id)
        if not props:
            continue

        partition_id = props.get("partitionId", "?")
        cursors[partition_id] = {
            "orchestration_id": props.get("orchestration_id", ""),
            "processed_count": props.get("processedCount", 0),
            "last_updated_ms": props.get("lastUpdated", 0),
        }

    return cursors


def _build_run_summary(
    orch: dict[str, Any],
    cursor_states: dict[str, dict[str, Any]],
    stale_cutoff: datetime,
) -> dict[str, Any]:
    """
    Build a run summary dict from stored orchestration state + live cursor data.

    Args:
        orch: Row from _list_orchestration_states()
        cursor_states: Cursor states for this orchestration's partitions (keyed by partitionId)
        stale_cutoff: Runs with last_updated before this are considered stale
    """
    completed_partitions: dict = orch.get("completed_partitions") or {}
    failed_partitions: dict = orch.get("failed_partitions") or {}
    partition_call_ids: dict = orch.get("partition_call_ids") or {}
    partition_retry_count: dict = orch.get("partition_retry_count") or {}
    total_partitions: int = orch.get("partition_count") or 0

    completed_ids = set(completed_partitions.keys())
    failed_ids = set(failed_partitions.keys())
    all_triggered_ids = set(partition_call_ids.keys())
    in_progress_ids = all_triggered_ids - completed_ids - failed_ids

    completed_count = len(completed_ids)
    failed_count = len(failed_ids)
    in_progress_count = len(in_progress_ids)

    # Age calculation
    last_updated_str: str | None = orch.get("last_updated")
    age_hours: float | None = None
    is_stale = False
    if last_updated_str:
        try:
            lu = datetime.fromisoformat(last_updated_str.replace("Z", "+00:00"))
            age_hours = (datetime.now(timezone.utc) - lu).total_seconds() / 3600.0
            is_stale = orch.get("status") in ("monitoring", "in_progress") and lu < stale_cutoff
        except (ValueError, TypeError):
            pass

    # Build per-partition detail
    partition_detail: dict[str, dict[str, Any]] = {}
    for pid in all_triggered_ids:
        if pid in completed_ids:
            entry: dict[str, Any] = {"status": "completed"}
        elif pid in failed_ids:
            retries = partition_retry_count.get(pid, 0)
            entry = {"status": "failed", "retry_count": retries}
        else:
            entry = {"status": "in_progress"}
            cursor = cursor_states.get(pid)
            if cursor:
                entry["processed_count"] = cursor.get("processed_count", 0)
                entry["last_updated_ms"] = cursor.get("last_updated_ms")

        partition_detail[pid] = entry

    # validated_instances in OrchestrationState only reflects completed partitions
    # (updated by the monitor schedule). Cursors are deleted on partition completion,
    # so in-progress processedCount + stored validated_instances = live total, no double-counting.
    validated = orch.get("validated_instances", 0)
    live_in_progress = sum(
        cursor_states[pid].get("processed_count", 0) for pid in in_progress_ids if pid in cursor_states
    )
    validated_live = validated + live_in_progress
    total = orch.get("total_instances", 0)
    violations = orch.get("total_violations", 0)
    pct = (validated_live / total * 100.0) if total > 0 else 0.0

    return {
        "orchestration_id": orch.get("orchestration_id", ""),
        "view_space": orch.get("view_space", ""),
        "view_external_id": orch.get("view_external_id", ""),
        "status": orch.get("status", ""),
        "partitions": {
            "total": total_partitions,
            "completed": completed_count,
            "failed": failed_count,
            "in_progress": in_progress_count,
        },
        "progress": {
            "validated_instances": validated_live,
            "total_instances": total,
            "total_violations": violations,
            "percent_complete": round(pct, 1),
        },
        "partition_detail": partition_detail,
        "created_time": orch.get("created_time"),
        "last_updated": last_updated_str,
        "age_hours": round(age_hours, 2) if age_hours is not None else None,
        "is_stale": is_stale,
    }
