"""QRP (qRaptor Routing Protocol) message models.

All message types exchanged between qRunX (server) and qRemoteX (client)
over the persistent WebSocket connection.  JSON-over-WebSocket, UTF-8.
"""

from __future__ import annotations

from enum import Enum
from typing import Any


# ──────────────────────────────────────────────── enums ──

class MessageType(str, Enum):
    """All QRP message types (both directions)."""

    # Client → Server
    HELLO = "hello"
    HEARTBEAT = "heartbeat"
    PONG = "pong"
    ACK = "ack"
    STATUS = "status"
    LOG = "log"
    RESULT = "result"
    RESULT_CHUNK_START = "result_chunk_start"
    RESULT_CHUNK = "result_chunk"
    RESULT_CHUNK_END = "result_chunk_end"
    ERROR = "error"
    UPLOAD_FILE = "upload_file"
    METRICS = "metrics"

    # Server → Client
    EXECUTE_SCRIPT = "execute_script"
    EXECUTE_TOOL = "execute_tool"
    CANCEL = "cancel"
    PING = "ping"
    DOWNLOAD_FILE_START = "download_file_start"
    DOWNLOAD_FILE_CHUNK = "download_file_chunk"
    DOWNLOAD_FILE_COMPLETE = "download_file_complete"
    CONFIG_UPDATE = "config_update"
    KEY_ROTATION_NOTICE = "key_rotation_notice"
    RESULT_CHUNK_RETRY = "result_chunk_retry"
    DRAIN = "drain"


class ExecutionStatus(str, Enum):
    QUEUED = "queued"
    STARTED = "started"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"
    TIMEOUT = "timeout"


class LogLevel(str, Enum):
    INFO = "info"
    WARN = "warn"
    ERROR = "error"
    STDOUT = "stdout"
    STDERR = "stderr"


class ExecutionTargetKind(str, Enum):
    NATIVE = "NATIVE"
    REMOTE_HTTP = "REMOTE_HTTP"


# ──────────────────────────────────────────────── helpers ──

def hello_message(
    executor_id: str,
    api_key_id: str,
    version: str,
    capabilities: list[str],
    labels: dict[str, str],
    system_info: dict[str, Any],
    *,
    name: str = "",
    max_concurrency: int = 25,
    group: str = "",
) -> dict[str, Any]:
    msg: dict[str, Any] = {
        "type": MessageType.HELLO.value,
        "executor_id": executor_id,
        "api_key_id": api_key_id,
        "version": version,
        "capabilities": capabilities,
        "labels": labels,
        "system_info": system_info,
        "max_concurrency": max_concurrency,
    }
    if name:
        msg["name"] = name
    if group:
        msg["group"] = group
    return msg


def heartbeat_message(
    executor_id: str,
    ts: int,
    metrics: dict[str, Any],
) -> dict[str, Any]:
    return {
        "type": MessageType.HEARTBEAT.value,
        "executor_id": executor_id,
        "ts": ts,
        "metrics": metrics,
    }


def pong_message(executor_id: str) -> dict[str, Any]:
    return {
        "type": MessageType.PONG.value,
        "executor_id": executor_id,
    }


def ack_message(request_id: str, executor_id: str) -> dict[str, Any]:
    return {
        "type": MessageType.ACK.value,
        "request_id": request_id,
        "executor_id": executor_id,
    }


def status_message(
    request_id: str,
    executor_id: str,
    status: ExecutionStatus,
    *,
    position: int | None = None,
) -> dict[str, Any]:
    msg: dict[str, Any] = {
        "type": MessageType.STATUS.value,
        "request_id": request_id,
        "executor_id": executor_id,
        "status": status.value,
    }
    if position is not None:
        msg["position"] = position
    return msg


def log_message(
    request_id: str,
    executor_id: str,
    level: str,
    stream: str,
    message: str,
    timestamp: str,
    *,
    run_id: str = "",
) -> dict[str, Any]:
    msg: dict[str, Any] = {
        "type": MessageType.LOG.value,
        "request_id": request_id,
        "executor_id": executor_id,
        "level": level,
        "stream": stream,
        "message": message,
        "timestamp": timestamp,
    }
    if run_id:
        msg["run_id"] = run_id
    return msg


def result_message(
    request_id: str,
    executor_id: str,
    success: bool,
    result: Any,
    duration_ms: int,
    *,
    peak_memory_mb: int | None = None,
    result_size_bytes: int | None = None,
    run_id: str = "",
    console_logs: list[str] | None = None,
) -> dict[str, Any]:
    msg: dict[str, Any] = {
        "type": MessageType.RESULT.value,
        "request_id": request_id,
        "executor_id": executor_id,
        "success": success,
        "result": result,
        "duration_ms": duration_ms,
    }
    if peak_memory_mb is not None:
        msg["peak_memory_mb"] = peak_memory_mb
    if result_size_bytes is not None:
        msg["result_size_bytes"] = result_size_bytes
    if run_id:
        msg["run_id"] = run_id
    if console_logs:
        msg["console_logs"] = console_logs
    return msg


def error_message(
    request_id: str,
    executor_id: str,
    error: str,
    *,
    trace: str | None = None,
    run_id: str = "",
) -> dict[str, Any]:
    msg: dict[str, Any] = {
        "type": MessageType.ERROR.value,
        "request_id": request_id,
        "executor_id": executor_id,
        "error": error,
    }
    if trace:
        msg["trace"] = trace
    if run_id:
        msg["run_id"] = run_id
    return msg


def result_chunk_start_message(
    request_id: str,
    total_chunks: int,
    total_size_bytes: int,
    content_type: str = "application/json",
    checksum_algo: str = "sha256",
) -> dict[str, Any]:
    return {
        "type": MessageType.RESULT_CHUNK_START.value,
        "request_id": request_id,
        "total_chunks": total_chunks,
        "total_size_bytes": total_size_bytes,
        "content_type": content_type,
        "checksum_algo": checksum_algo,
    }


def result_chunk_message(
    request_id: str,
    chunk_index: int,
    data_b64: str,
) -> dict[str, Any]:
    return {
        "type": MessageType.RESULT_CHUNK.value,
        "request_id": request_id,
        "chunk_index": chunk_index,
        "data_b64": data_b64,
    }


def result_chunk_end_message(
    request_id: str,
    success: bool,
    duration_ms: int,
    checksum: str,
) -> dict[str, Any]:
    return {
        "type": MessageType.RESULT_CHUNK_END.value,
        "request_id": request_id,
        "success": success,
        "duration_ms": duration_ms,
        "checksum": checksum,
    }
