# Stage 1: Build patched Rust FFI with HKDF support for Element Call E2EE
# Fork: onestacked/livekit-rust-sdks branch EC-compat-changes
# PR: https://github.com/livekit/rust-sdks/pull/904
FROM rust:1.85-slim-bookworm AS rust-build
RUN apt-get update && apt-get install -y --no-install-recommends \
    git cmake g++ libssl-dev pkg-config protobuf-compiler \
    && rm -rf /var/lib/apt/lists/*
WORKDIR /build
RUN git clone --branch EC-compat-changes --depth 1 \
    https://github.com/onestacked/livekit-rust-sdks.git
WORKDIR /build/livekit-rust-sdks/livekit-ffi
RUN cargo build --release

# Stage 2: Generate Python protobuf bindings from patched .proto files
FROM python:3.11-slim-bookworm AS proto-gen
RUN pip install --no-cache-dir protobuf grpcio-tools mypy-protobuf
COPY --from=rust-build /build/livekit-rust-sdks/livekit-ffi/protocol/ /proto/
RUN mkdir -p /gen && \
    python -m grpc_tools.protoc \
      -I/proto \
      --python_out=/gen \
      --mypy_out=/gen \
      /proto/audio_frame.proto \
      /proto/ffi.proto \
      /proto/handle.proto \
      /proto/participant.proto \
      /proto/room.proto \
      /proto/track.proto \
      /proto/video_frame.proto \
      /proto/e2ee.proto \
      /proto/stats.proto \
      /proto/track_publication.proto \
      /proto/rpc.proto \
      /proto/data_stream.proto && \
    touch /gen/__init__.py && \
    # Fix imports to be relative (same as upstream generate_proto.sh)
    for f in /gen/*.py /gen/*.pyi; do \
      perl -i -pe 's|^(import (audio_frame_pb2\|ffi_pb2\|handle_pb2\|participant_pb2\|room_pb2\|track_pb2\|video_frame_pb2\|e2ee_pb2\|stats_pb2\|rpc_pb2\|track_publication_pb2\|data_stream_pb2))|from . \1|g' "$f"; \
    done

# Stage 3: Final image
FROM python:3.11-slim-bookworm
WORKDIR /app
RUN apt-get update && apt-get install -y --no-install-recommends \
    ffmpeg libolm-dev \
    && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Overwrite installed FFI binary with patched version (HKDF + key_ring_size support)
COPY --from=rust-build /build/livekit-rust-sdks/target/release/liblivekit_ffi.so /patched/
ENV LIVEKIT_LIB_PATH=/patched/liblivekit_ffi.so

# Overwrite installed proto bindings with patched versions (new fields: key_ring_size, key_derivation_function)
COPY --from=proto-gen /gen/ /patched_proto/
RUN PROTO_DIR=$(python -c "import livekit.rtc._proto; import os; print(os.path.dirname(livekit.rtc._proto.__file__))") && \
    cp /patched_proto/*.py "$PROTO_DIR/" && \
    cp /patched_proto/*.pyi "$PROTO_DIR/" 2>/dev/null || true

# Patch SDK Python code to pass new fields through to proto (e2ee.py + room.py)
COPY patch_sdk.py /tmp/patch_sdk.py
RUN python /tmp/patch_sdk.py && rm /tmp/patch_sdk.py

COPY . .