#!/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 = "./test.zip"

# 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="test.zip"
    )
    
    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)")
    log_trace("Received message:\n$msg")
    
    # Use SmartReceive to handle the data
    result = SmartReceive(msg)
    
    # Check transport type
    if result.envelope.transport == "direct"
        log_trace("Received direct transport with $(length(result.data)) bytes")
    else
        # For link transport, result.data is the URL
        log_trace("Received link transport at $(result.data)")
    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
        output_path = "test_output.bin"
        write(output_path, result.data)
        log_trace("Saved received data to $output_path")
        
        # Verify file size
        original_size = length(read(LARGE_FILE_PATH))
        if file_size == original_size
            log_trace("SUCCESS: File size matches! Original: $original_size bytes")
        else
            log_trace("WARNING: File size mismatch! Original: $original_size, Received: $file_size")
        end
      end
    end
  end
        
  # Keep listening for 10 seconds
  sleep(10)  
  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.")