#!/usr/bin/env julia
# Test script for large payload testing using binary transport
# Tests sending a large file (> 1MB) via smartsend with binary type

using NATS, JSON, UUIDs, Dates

# Include the bridge module
include("../src/NATSBridge.jl")
using .NATSBridge

# Configuration
const SUBJECT = "/large_binary_test"
const NATS_URL = "nats.yiem.cc"
const FILESERVER_URL = "http://192.168.88.104:8080"

# Create correlation ID for tracing
correlation_id = string(uuid4())

# File path for large binary payload test
const LARGE_FILE_PATH = "./testFile_small.zip"
const filename = basename(LARGE_FILE_PATH)

# Helper: Log with correlation ID
function log_trace(message)
    timestamp = Dates.now()
    println("[$timestamp] [Correlation: $correlation_id] $message")
end

# Sender: Send large binary file via smartsend
function test_large_binary_send()
    conn = NATS.connect(NATS_URL)
    # Read the large file as binary data
    log_trace("Reading large file: $LARGE_FILE_PATH")
    file_data = read(LARGE_FILE_PATH)
    
    file_size = length(file_data)
    log_trace("File size: $file_size bytes")
    
    # Use smartsend with binary type - will automatically use link transport
    # if file size exceeds the threshold (1MB by default)
    env = NATSBridge.smartsend(
        SUBJECT,
        file_data,
        "binary",
        nats_url = NATS_URL,
        fileserver_url = FILESERVER_URL;
        dataname=filename
    )
    
    log_trace("Sent message with transport: $(env.transport)")
    log_trace("Envelope type: $(env.type)")
    
    # Check if link transport was used
    if env.transport == "link"
        log_trace("Using link transport - file uploaded to HTTP server")
        log_trace("URL: $(env.url)")
    else
        log_trace("Using direct transport - payload sent via NATS")
    end
    
    NATS.drain(conn)
end

# Receiver: Listen for messages and verify large payload handling
function test_large_binary_receive()
  conn = NATS.connect(NATS_URL)
  NATS.subscribe(conn, SUBJECT) do msg
    log_trace("Received message on $(msg.subject)")
    
    # Use NATSBridge.smartreceive to handle the data
    result = NATSBridge.smartreceive(msg)
    # Check transport type
    if result.envelope.transport == "direct"
        log_trace("Received direct transport")
    else
        # For link transport, result.data is the URL
        log_trace("Received link transport")
    end
              
    # Verify the received data matches the original
    if result.envelope.type == "binary"
      if isa(result.data, Vector{UInt8})
        file_size = length(result.data)
        log_trace("Received $(file_size) bytes of binary data")
        
        # Save received data to a test file
        println("metadata ", result.envelope.metadata)
        dataname = result.envelope.metadata["dataname"]
        if dataname != "NA"
          output_path = "./new_$dataname"
          write(output_path, result.data)
          log_trace("Saved received data to $output_path")
        end
        
        # Verify file size
        original_size = length(read(LARGE_FILE_PATH))
        if file_size == result.envelope.metadata["content_length"]
            log_trace("SUCCESS: File size matches! Original: $(result.envelope.metadata["content_length"]) bytes")
        else
            log_trace("WARNING: File size mismatch! Original: $(result.envelope.metadata["content_length"]), Received: $file_size")
        end
      end
    end
  end

  # Keep listening for 10 seconds
  sleep(120)  
  NATS.drain(conn)
end

# Run the test
println("Starting large binary payload test...")
println("Correlation ID: $correlation_id")
println("Large file: $LARGE_FILE_PATH")

# # Run sender first
# println("start smartsend")
# test_large_binary_send()

# Run receiver
println("testing smartreceive")
test_large_binary_receive()

println("Test completed.")