"""Celery event capture and normalization.

Captures Celery task and worker events and normalizes them into
framework-agnostic JSON payloads matching the Sluice ingest API schema
(``ingestEventSchema`` in ``dashboard/lib/schemas/api.ts``).

Wire format: nested camelCase with ISO 8601 timestamps.

State mapping follows docs/architecture/multi_framework_vision.md §2.3:

    Celery PENDING  -> unknown  (no information)
    Celery RECEIVED -> queued
    Celery STARTED  -> active
    Celery SUCCESS  -> completed
    Celery FAILURE  -> failed
    Celery RETRY    -> retrying
    Celery REVOKED  -> cancelled
    Celery REJECTED -> failed
"""

from __future__ import annotations

from datetime import UTC, datetime, timedelta
from typing import Any

# Celery state -> Sluice unified state.
_CELERY_STATE_MAP: dict[str, str] = {
    "PENDING": "unknown",
    "RECEIVED": "queued",
    "STARTED": "active",
    "SUCCESS": "completed",
    "FAILURE": "failed",
    "RETRY": "retrying",
    "REVOKED": "cancelled",
    "REJECTED": "failed",
}

# Celery task event type -> Celery state string.
_EVENT_TYPE_TO_STATE: dict[str, str] = {
    "task-sent": "PENDING",
    "task-received": "RECEIVED",
    "task-started": "STARTED",
    "task-succeeded": "SUCCESS",
    "task-failed": "FAILURE",
    "task-retried": "RETRY",
    "task-revoked": "REVOKED",
    "task-rejected": "REJECTED",
}

# Worker event types that we capture.
_WORKER_EVENT_TYPES: frozenset[str] = frozenset(
    {
        "worker-online",
        "worker-heartbeat",
        "worker-offline",
    }
)

_WORKER_STATE_MAP: dict[str, str] = {
    "worker-online": "online",
    "worker-heartbeat": "heartbeat",
    "worker-offline": "offline",
}


def _to_iso8601(ts: float) -> str:
    """Convert a Unix timestamp to an ISO 8601 datetime string with Z suffix.

    Zod's ``z.string().datetime()`` requires the ``Z`` suffix — it rejects
    Python's default ``+00:00`` format.
    """
    return datetime.fromtimestamp(ts, tz=UTC).isoformat().replace("+00:00", "Z")


def _now_iso8601() -> str:
    """Return the current time as an ISO 8601 datetime string with Z suffix."""
    return datetime.now(tz=UTC).isoformat().replace("+00:00", "Z")


def normalize_task_event(event: dict[str, Any]) -> dict[str, Any] | None:
    """Convert a raw Celery task event into a normalized Sluice ingest event.

    Output matches ``ingestEventSchema`` from the dashboard API schema:
    - ``type``: event type string (e.g. "job.state_change")
    - ``timestamp``: ISO 8601 datetime string
    - ``framework``: "celery"
    - ``job``: nested object with camelCase keys

    ``connectionId`` is NOT included here — it is injected by the transport
    layer which holds the connection context.

    Returns None if the event type is unknown or the event is missing
    required fields (uuid).
    """
    event_type: str = event.get("type", "")
    celery_state = _EVENT_TYPE_TO_STATE.get(event_type)
    if celery_state is None:
        return None

    uuid: str = event.get("uuid", "")
    if not uuid:
        return None

    unified_state = _CELERY_STATE_MAP[celery_state]

    # Convert timestamp to ISO 8601.
    raw_ts = event.get("timestamp")
    try:
        timestamp = _to_iso8601(float(raw_ts)) if raw_ts else _now_iso8601()
    except (ValueError, TypeError):
        timestamp = _now_iso8601()

    # Build the nested job object (camelCase keys per ingestEventSchema).
    job: dict[str, Any] = {
        "externalId": uuid,
        "name": event.get("name", "") or "unknown",
        "state": unified_state,
    }

    # Optional job fields.
    queue = event.get("queue", "")
    if queue:
        job["queue"] = queue

    worker_id = event.get("hostname", "")
    if worker_id:
        job["workerId"] = worker_id

    # Retry count -> attempt (1-based).
    if "retries" in event:
        job["attempt"] = int(event["retries"]) + 1

    # Error information (task-failed, task-rejected).
    if "exception" in event:
        job["error"] = str(event["exception"])
    if "traceback" in event:
        job["stacktrace"] = str(event["traceback"])

    # Runtime -> completedAt/startedAt timing.
    # Celery's task-succeeded event includes `runtime` (seconds).
    # The server computes durationMs from startedAt/completedAt.
    if celery_state == "SUCCESS" and "runtime" in event:
        try:
            runtime_sec = float(event["runtime"])
            # completedAt = event timestamp, startedAt = completedAt - runtime
            completed_dt = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
            started_dt = completed_dt - timedelta(seconds=runtime_sec)
            job["completedAt"] = timestamp
            job["startedAt"] = started_dt.isoformat().replace("+00:00", "Z")
        except (ValueError, TypeError, OverflowError):
            pass

    # Extensions: framework-specific fields.
    extensions: dict[str, Any] = {}
    if "eta" in event:
        extensions["eta"] = event["eta"]
    if "expires" in event:
        extensions["expires"] = event["expires"]
    if "exchange" in event:
        extensions["exchange"] = event["exchange"]
    if "routing_key" in event:
        extensions["routing_key"] = event["routing_key"]
    # Preserve original Celery state for debugging.
    extensions["celeryState"] = celery_state
    if extensions:
        job["extensions"] = extensions

    return {
        "type": "job.state_change",
        "timestamp": timestamp,
        "framework": "celery",
        "job": job,
    }


def normalize_worker_event(event: dict[str, Any]) -> dict[str, Any] | None:
    """Convert a raw Celery worker event into a normalized Sluice ingest event.

    Output matches ``ingestEventSchema`` with the ``worker`` nested object.

    ``connectionId`` is NOT included here — it is injected by the transport.

    Returns None if the event type is unrecognized.
    """
    event_type: str = event.get("type", "")
    if event_type not in _WORKER_EVENT_TYPES:
        return None

    hostname: str = event.get("hostname", "")
    if not hostname:
        return None

    worker_state = _WORKER_STATE_MAP[event_type]

    # Convert timestamp to ISO 8601.
    raw_ts = event.get("timestamp")
    try:
        timestamp = _to_iso8601(float(raw_ts)) if raw_ts else _now_iso8601()
    except (ValueError, TypeError):
        timestamp = _now_iso8601()

    # Build the nested worker object (camelCase keys per ingestEventSchema).
    worker: dict[str, Any] = {
        "externalId": hostname,
        "hostname": hostname,
        "state": worker_state,
    }

    # Optional worker fields.
    if "pid" in event:
        worker["pid"] = event["pid"]
    if "active" in event:
        worker["activeJobs"] = event["active"]
    if event.get("concurrency"):
        worker["concurrency"] = event["concurrency"]

    # Worker metadata (non-schema fields go into metadata).
    metadata: dict[str, Any] = {}
    if "sw_ident" in event:
        metadata["software"] = event["sw_ident"]
    if "sw_ver" in event:
        metadata["softwareVersion"] = event["sw_ver"]
    if "processed" in event:
        metadata["processedTotal"] = event["processed"]
    if "loadavg" in event:
        metadata["loadavg"] = event["loadavg"]
    if "freq" in event:
        metadata["heartbeatFreq"] = event["freq"]
    if metadata:
        worker["metadata"] = metadata

    return {
        "type": "worker.state_change",
        "timestamp": timestamp,
        "framework": "celery",
        "worker": worker,
    }
