#!/usr/bin/env python3
"""
Ray transport for the VCN-LUPINE bridge.

Replaces GPUNodeConnection's TCP/MKV transport with Ray ObjectRef transport.
FrameDispatcher, backends, and frame protocol are unchanged — only the
wire between daemon and GPU node changes.

Usage:
    # As a drop-in replacement for GPUNodeConnection:
    from ray_vcn_bridge import RayGPUNodeConnection
    gpu = RayGPUNodeConnection()
    reply = gpu.send_frame(TAG_STRAND, payload)

    # As a standalone Ray actor:
    import ray
    ray.init("auto")
    bridge = RayVCNBridge.remote()
    reply_ref = bridge.dispatch_frame.remote(TAG_STRAND, payload)
    reply = ray.get(reply_ref)

Design:
    - Existing FrameDispatcher handles all tag routing
    - Existing BraidBackend / CUDABackend handle compute
    - Ray replaces: Unix socket IPC + TCP transport + MKV encode/decode
    - ObjectRef replaces: MKV bytes over TCP (zero-copy on same node)
    - FAMM gate check happens BEFORE ray.put() (same as before send_frame)
"""

from __future__ import annotations

import sys
import time
from pathlib import Path
from typing import Optional

import ray

# VCN bridge imports (unchanged — these are the abstraction layer)
sys.path.insert(0, str(Path(__file__).resolve().parent))
from vcn_lupine_bridge import (
    TAG_STRAND, TAG_CROSSING, TAG_PIST, TAG_LUPINE, TAG_VAAPI,
    FLAG_REPLY, FRAME_HDR_SIZE,
    pack_frame, unpack_frame, pack_reply, unpack_reply,
    tag_name, bridge_receipt,
    FrameDispatcher, CUDABackend, BraidBackend, VAAPIBackend,
)


@ray.remote(runtime_env={"pip": ["reedsolo", "cryptography"]})
class RayVAAPIBackend(VAAPIBackend):
    """VA-API compute backend running as a Ray actor.

    Uses AMD/Intel VAAPI for hardware H.264 encode/decode.
    Scheduled on nodes with AMD iGPU or Intel GPU.
    """

    def __init__(self, device: str = "/dev/dri/renderD128", encoder: str = "h264_vaapi"):
        self.device = device
        self.encoder = encoder

    def encode(self, payload: bytes, resolution: str = "1080p") -> bytes:
        from braid_vcn_encoder import encode_braid_strand
        import json
        strand = json.loads(payload.decode()) if payload else {}
        return encode_braid_strand(strand, resolution)

    def decode(self, mkv_data: bytes) -> bytes:
        from braid_vcn_encoder import decode_braid_mkv
        result = decode_braid_mkv(mkv_data)
        return json.dumps(result, default=str).encode()


# ── Sync wrappers for Ray actors (FrameDispatcher expects sync interface) ──

class SyncBraidWrapper(BraidBackend):
    """Wraps a Ray BraidBackend actor to satisfy FrameDispatcher's sync interface."""
    def __init__(self, actor):
        self._actor = actor
    def compute(self, tag: int, payload: bytes) -> bytes:
        return ray.get(self._actor.compute.remote(tag, payload))


class SyncCUDAWrapper(CUDABackend):
    """Wraps a Ray CUDABackend actor to satisfy FrameDispatcher's sync interface."""
    def __init__(self, actor):
        self._actor = actor
    def call(self, opcode: int, args: dict):
        return ray.get(self._actor.call.remote(opcode, args))


class SyncVAAPIWrapper(VAAPIBackend):
    """Wraps a Ray VAAPIBackend actor to satisfy FrameDispatcher's sync interface."""
    def __init__(self, actor):
        self._actor = actor
    def encode(self, payload: bytes, resolution: str = "1080p") -> bytes:
        return ray.get(self._actor.encode.remote(payload, resolution))
    def decode(self, mkv_data: bytes) -> bytes:
        return ray.get(self._actor.decode.remote(mkv_data))


# ── Ray-backed BraidBackend ─────────────────────────────────────────────────

@ray.remote(runtime_env={"pip": ["reedsolo", "cryptography"]})
class RayBraidBackend(BraidBackend):
    """Braid compute backend running as a Ray actor.

    Delegates to vcn_compute_substrate for actual computation.
    Scheduled on workers with /dev/dri access for hardware encode.
    """

    def compute(self, tag: int, payload: bytes) -> bytes:
        """Compute braid operation — matches VCNBraidBackend from vcn_lupine_gpu_node."""
        from vcn_compute_substrate import (
            TAG_STRAND, TAG_CROSSING, TAG_PIST,
            deserialize_braid_strand, deserialize_braid_bracket,
            compute_strand_phase, compute_crossing_residual,
            serialize_crossing_result,
            BRAID_BRACKET_BYTES,
        )
        import json

        if tag == TAG_STRAND:
            strand = deserialize_braid_strand(payload)
            result = compute_strand_phase(strand)
            return json.dumps(result, default=str).encode()

        elif tag == TAG_CROSSING:
            bracket_a = deserialize_braid_bracket(payload[:BRAID_BRACKET_BYTES])
            bracket_b = deserialize_braid_bracket(payload[BRAID_BRACKET_BYTES:2*BRAID_BRACKET_BYTES])
            result = compute_crossing_residual(bracket_a, bracket_b)
            return serialize_crossing_result(result)

        elif tag == TAG_PIST:
            spectral_data = json.loads(payload.decode("utf-8"))
            return json.dumps({"status": "ok", "input_keys": list(spectral_data.keys())}).encode()

        return b""


# ── Ray-backed CUDABackend ──────────────────────────────────────────────────

@ray.remote(num_gpus=1, runtime_env={"pip": ["reedsolo", "cryptography"]})
class RayCUDABackend(CUDABackend):
    """CUDA compute backend running as a Ray actor on GPU worker.

    Uses /dev/dri directly (Mesa/NVK) instead of NVIDIA container runtime.
    """

    def call(self, opcode: int, args: dict):
        """Execute CUDA operation via Vulkan compute or libcuda."""
        # For now, delegate to the LUPINE HTTP endpoint
        # Future: direct Vulkan compute via NVK
        from vcn_lupine_bridge import LUPINEBackend
        backend = LUPINEBackend()
        return backend.call(opcode, args)


# ── Ray GPU Node Connection (drop-in for GPUNodeConnection) ─────────────────

class RayGPUNodeConnection:
    """Drop-in replacement for GPUNodeConnection using Ray transport.

    Same interface: send_frame(tag, payload) → reply_bytes
    Same interface: send_frame_async(tag, payload)
    Different transport: Ray ObjectRef instead of TCP/MKV
    """

    def __init__(self, dispatcher: Optional[FrameDispatcher] = None):
        self._seq = 0
        self._dispatcher = dispatcher
        self._braid_backend = None
        self._cuda_backend = None

    def connect(self):
        """Initialize Ray backends (replaces TCP connect)."""
        if not ray.is_initialized():
            ray.init("auto", ignore_reinit_error=True)

        self._braid_actor = RayBraidBackend.remote()
        self._cuda_actor = RayCUDABackend.remote()
        self._braid_backend = self._braid_actor
        self._cuda_backend = self._cuda_actor

        if self._dispatcher is None:
            self._dispatcher = FrameDispatcher(
                cuda_backend=SyncCUDAWrapper(self._cuda_actor),
                braid_backend=SyncBraidWrapper(self._braid_actor),
            )

    def send_frame(self, tag: int, payload: bytes, timeout: float = 30.0) -> bytes:
        """Send a frame via Ray and wait for reply.

        Replaces: pack_frame → MKV encode → TCP send → wait → MKV decode
        With:     FrameDispatcher.dispatch() on Ray actor (zero-copy)
        """
        self._seq += 1
        seq = self._seq

        receipt = bridge_receipt(tag, 0, seq, len(payload), True)

        if tag == TAG_LUPINE:
            # LUPINE calls go to GPU worker
            reply_ref = self._cuda_backend.call.remote(
                *self._decode_lupine_args(payload))
            result = ray.get(reply_ref, timeout=timeout)
            from vcn_lupine_bridge import encode_lupine_reply
            reply_payload = encode_lupine_reply(0, 0, result)
            return pack_reply(TAG_LUPINE, seq, reply_payload)
        else:
            # Braid ops go to compute worker
            reply_ref = self._braid_backend.compute.remote(tag, payload)
            result = ray.get(reply_ref, timeout=timeout)
            return pack_reply(tag, seq, result)

    def send_frame_async(self, tag: int, payload: bytes):
        """Fire-and-forget frame dispatch (no reply expected)."""
        self._seq += 1
        if tag == TAG_LUPINE:
            self._cuda_backend.call.remote(
                *self._decode_lupine_args(payload))
        else:
            self._braid_backend.compute.remote(tag, payload)

    def close(self):
        """Clean up Ray actors."""
        if self._braid_backend:
            ray.kill(self._braid_backend)
        if self._cuda_backend:
            ray.kill(self._cuda_backend)

    @staticmethod
    def _decode_lupine_args(payload: bytes):
        """Extract opcode and args from LUPINE payload."""
        from vcn_lupine_bridge import decode_lupine_request
        request_id, opcode, args = decode_lupine_request(payload)
        return opcode, args


# ── Ray VCN Bridge actor (full bridge as Ray actor) ─────────────────────────

@ray.remote(runtime_env={"pip": ["reedsolo", "cryptography"]})
class RayVCNBridge:
    """VCN-LUPINE bridge as a Ray actor.

    Replaces the Unix socket daemon + TCP connection.
    Producers call dispatch_frame() instead of write_ipc_frame().
    """

    def __init__(self):
        self._braid_actor = RayBraidBackend.remote()
        self._cuda_actor = RayCUDABackend.remote()
        self._dispatcher = FrameDispatcher(
            cuda_backend=SyncCUDAWrapper(self._cuda_actor),
            braid_backend=SyncBraidWrapper(self._braid_actor),
        )
        self._braid = self._braid_actor
        self._cuda = self._cuda_actor
        self._seq = 0
        self._receipts = []

    def dispatch_frame(self, tag: int, payload: bytes) -> bytes:
        """Dispatch a frame through the FrameDispatcher.

        Same semantics as daemon's handle_ipc_client: receives (tag, payload),
        dispatches to appropriate backend, returns reply frame.
        """
        self._seq += 1
        seq = self._seq

        receipt = bridge_receipt(tag, 0, seq, len(payload), True)
        self._receipts.append(receipt)

        reply = self._dispatcher.dispatch(tag, 0, seq, payload)
        return reply if reply else b""

    def dispatch_frame_ref(self, tag: int, ref: ray.ObjectRef) -> bytes:
        """Dispatch from an ObjectRef (zero-copy on same node)."""
        payload = ray.get(ref)
        return self.dispatch_frame(tag, payload)

    def dispatch_batch(self, frames: list[tuple[int, bytes]]) -> list[bytes]:
        """Dispatch multiple frames in parallel."""
        refs = [self.dispatch_frame.remote(tag, payload)
                for tag, payload in frames]
        return ray.get(refs)

    def get_receipts(self) -> list[dict]:
        """Return all dispatch receipts."""
        return self._receipts

    def get_stats(self) -> dict:
        """Return bridge statistics."""
        return {
            "frames_dispatched": self._seq,
            "receipts": len(self._receipts),
            "braid_backend": self._braid is not None,
            "cuda_backend": self._cuda is not None,
        }


# ── Convenience: daemon replacement ─────────────────────────────────────────

def create_ray_daemon() -> RayVCNBridge:
    """Create a Ray VCN bridge actor (replaces daemon startup).

    Returns an actor handle that accepts dispatch_frame() calls.
    """
    if not ray.is_initialized():
        ray.init("auto", ignore_reinit_error=True)
    return RayVCNBridge.options(name="vcn-bridge", lifetime="detached").remote()


# ── CLI ──────────────────────────────────────────────────────────────────────

def main():
    """Smoke test for the Ray VCN bridge."""
    import json

    print("Ray VCN Bridge — connecting...")
    ray.init("auto", ignore_reinit_error=True)

    bridge = RayVCNBridge.remote()

    # Test strand dispatch
    from braid_vcn_encoder import _serialize_strand
    test_strand = {
        "phaseAcc": {"x": 42, "y": 100},
        "parity": True, "slot": 7, "residue": 3, "jitter": 0,
        "bracket": {"lower": 10000, "upper": 20000, "gap": 5000,
                     "kappa": 32768, "phi": 16384, "admissible": True},
    }
    payload = _serialize_strand(test_strand)

    print(f"Dispatching TAG_STRAND ({len(payload)} bytes)...")
    reply = ray.get(bridge.dispatch_frame.remote(TAG_STRAND, payload))
    print(f"  Reply: {len(reply)} bytes")

    stats = ray.get(bridge.get_stats.remote())
    print(f"  Stats: {json.dumps(stats, indent=2)}")

    # Test RayGPUNodeConnection drop-in
    print("\nTesting RayGPUNodeConnection...")
    conn = RayGPUNodeConnection()
    conn.connect()
    reply2 = conn.send_frame(TAG_STRAND, payload)
    print(f"  Reply via conn: {len(reply2)} bytes")
    conn.close()

    print("\n✓ Ray VCN Bridge smoke test passed")
    ray.shutdown()


if __name__ == "__main__":
    main()