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

using NATS, JSON, UUIDs, Dates
using HTTP

# 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())


# ------------------------------------------------------------------------------------------------ #
#                                        test file transfer                                        #
# ------------------------------------------------------------------------------------------------ #

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

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

# File upload handler for plik server
function plik_upload_handler(fileserver_url::String, dataname::String, data::Vector{UInt8})::Dict{String, Any}
    # Get upload ID
    url_getUploadID = "$fileserver_url/upload"
    headers = ["Content-Type" => "application/json"]
    body = """{ "OneShot" : true }"""
    httpResponse = HTTP.request("POST", url_getUploadID, headers, body; body_is_form=false)
    responseJson = JSON.parse(String(httpResponse.body))
    uploadid = responseJson["id"]
    uploadtoken = responseJson["uploadToken"]
    
    # Upload file
    file_multipart = HTTP.Multipart(dataname, IOBuffer(data), "application/octet-stream")
    url_upload = "$fileserver_url/file/$uploadid"
    headers = ["X-UploadToken" => uploadtoken]
    
    form = HTTP.Form(Dict("file" => file_multipart))
    httpResponse = HTTP.post(url_upload, headers, form)
    responseJson = JSON.parse(String(httpResponse.body))
    
    fileid = responseJson["id"]
    url = "$fileserver_url/file/$uploadid/$fileid/$dataname"
    
    return Dict("status" => httpResponse.status, "uploadid" => uploadid, "fileid" => fileid, "url" => url)
end

# File download handler for plik server
function plik_download_handler(fileserver_url::String, url::String, max_retries::Int, base_delay::Int, max_delay::Int)::Vector{UInt8}
    delay = base_delay
    for attempt in 1:max_retries
        try
            response = HTTP.request("GET", url)
            if response.status == 200
                return response.body
            else
                error("Failed to fetch: $(response.status)")
            end
        catch e
            if attempt < max_retries
                sleep(delay / 1000.0)
                delay = min(delay * 2, max_delay)
            end
        end
    end
    error("Failed to fetch data after $max_retries attempts")
end

# Sender: Send large binary file via smartsend
function test_large_binary_send()
    # Read the large file as binary data
    log_trace("Reading large file: $FILE_PATH")
    file_data = read(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)
    # API: smartsend(subject, [(dataname, data, type), ...]; keywords...)
    env = NATSBridge.smartsend(
        SUBJECT,
        [(filename, file_data, "binary")],  # List of (dataname, data, type) tuples
        nats_url = NATS_URL,
        fileserverUploadHandler = plik_upload_handler,
        size_threshold = 1_000_000,
        correlation_id = correlation_id,
        msg_purpose = "chat",
        sender_name = "sender",
        receiver_name = "",
        receiver_id = "",
        reply_to = "",
        reply_to_msg_id = ""
    )
    
    log_trace("Sent message with transport: $(env.payloads[1].transport)")
    log_trace("Envelope type: $(env.payloads[1].type)")
    
    # Check if link transport was used
    if env.payloads[1].transport == "link"
        log_trace("Using link transport - file uploaded to HTTP server")
        log_trace("URL: $(env.payloads[1].data)")
    else
        log_trace("Using direct transport - payload sent via NATS")
    end
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
        # API: smartreceive(msg, download_handler; max_retries, base_delay, max_delay)
        result = NATSBridge.smartreceive(
            msg,
            plik_download_handler,
            max_retries = 5,
            base_delay = 100,
            max_delay = 5000
        )
        
        # Result is a list of (dataname, data) tuples
        for (dataname, data) in result
            # Check transport type from the envelope
            # For link transport, data is the URL string
            # For direct transport, data is the actual payload bytes
            
            if isa(data, Vector{UInt8})
                file_size = length(data)
                log_trace("Received $(file_size) bytes of binary data for '$dataname'")
                
                # Save received data to a test file
                output_path = "./new_$dataname"
                write(output_path, data)
                log_trace("Saved received data to $output_path")
            else
                log_trace("Received data for '$dataname' (type: $(typeof(data)))")
            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("File: $FILE_PATH")

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

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

println("Test completed.")