import asyncio
import json
import base64

CHUNK_SIZE = 256 * 1024  # 256 KB

async def send_file_in_chunks(room, filepath, target_participant_sid=None):
    """
    room: LiveKit Room instance for the agent
    filepath: path to local file to send
    target_participant_sid: if set, send to that participant only (include in payload)
    """
    local_participant = room.local_participant

    # Read and stream file in chunks
    total_bytes = 0
    with open(filepath, "rb") as f:
        f.seek(0, 2)
        total_bytes = f.tell()
        f.seek(0)

        chunk_index = 0
        while True:
            chunk = f.read(CHUNK_SIZE)
            if not chunk:
                break

            payload = {
                "type": "bytestream_chunk",
                "file_name": filepath.split("/")[-1],
                "total_size": total_bytes,
                "chunk_index": chunk_index,
                "chunk_size": len(chunk),
                "target": target_participant_sid,
            }

            # Encode chunk as base64 to keep payload JSON-friendly (or send binaryover data channel directly if supported)
            payload["data_b64"] = base64.b64encode(chunk).decode("ascii")

            # You can choose reliability/reliability=RELIABLE for guaranteed delivery
            await local_participant.publish_data(
                data=json.dumps(payload).encode("utf-8"),
                reliability=1  # DataPublishReliability.RELIABLE equivalent
            )

            chunk_index += 1
            # Optional: throttle to avoid saturating network
            await asyncio.sleep(0)

    # Send finalization message
    finished_msg = {
        "type": "bytestream_finished",
        "file_name": filepath.split("/")[-1],
        "total_chunks": chunk_index,
        "total_size": total_bytes,
        "target": target_participant_sid,
    }
    await local_participant.publish_data(
        data=json.dumps(finished_msg).encode("utf-8"),
        reliability=1
    )

// room: connected LiveKit Room instance in the browser
const CHUNKS = {}; // store chunks per file_name

room.on("dataReceived", (payload, participant) => {
  // payload is an ArrayBuffer/Uint8Array depending on client lib; ensure string decode
  let text;
  try {
    text = new TextDecoder().decode(payload);
  } catch (e) {
    console.error("Failed to decode payload", e);
    return;
  }

  let msg;
  try {
    msg = JSON.parse(text);
  } catch (e) {
    console.error("Invalid JSON payload", e);
    return;
  }

  if (msg.type === "bytestream_chunk") {
    const name = msg.file_name;
    const chunkBytes = Uint8Array.from(atob(msg.data_b64), c => c.charCodeAt(0));

    if (!CHUNKS[name]) {
      CHUNKS[name] = {
        total_size: msg.total_size,
        received_size: 0,
        total_chunks: null,
        buffers: [],
      };
    }

    CHUNKS[name].buffers[msg.chunk_index] = chunkBytes;
    CHUNKS[name].received_size += msg.chunk_size;

    // Progress callback (percentage)
    const progress = (CHUNKS[name].received_size / CHUNKS[name].total_size) * 100;
    console.log(`Receiving ${name}: ${progress.toFixed(1)}%`);

  } else if (msg.type === "bytestream_finished") {
    const name = msg.file_name;
    const entry = CHUNKS[name];
    if (!entry) {
      console.error("No chunks for finished file", name);
      return;
    }

    // Concatenate chunks in order
    const buffers = entry.buffers;
    let totalLen = 0;
    for (const b of buffers) totalLen += b.length;
    const tmp = new Uint8Array(totalLen);
    let offset = 0;
    for (let i = 0; i < buffers.length; i++) {
      const b = buffers[i];
      tmp.set(b, offset);
      offset += b.length;
    }

    // Create blob and download or process
    const blob = new Blob([tmp]);
    const url = URL.createObjectURL(blob);
    const a = document.createElement("a");
    a.href = url;
    a.download = name;
    document.body.appendChild(a);
    a.click();
    a.remove();
    URL.revokeObjectURL(url);

    console.log(`File ${name} received and saved.`);
    delete CHUNKS[name];
  }
});