#!/usr/bin/env julia
# Test script for Dictionary transport testing
# Tests sending 1 large and 1 small Dictionaries via direct and link transport
# Uses NATSBridge.jl smartsend with "dictionary" type

using NATS, JSON, UUIDs, Dates, PrettyPrinting, DataFrames, Arrow, HTTP

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

# Configuration
const SUBJECT = "/NATSBridge_dict_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 dictionary transfer                                  #
# ------------------------------------------------------------------------------------------------ #


# 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


# Sender: Send Dictionaries via smartsend
function test_dict_send()
    # Create a small Dictionary (will use direct transport)
    small_dict = Dict(
        "name" => "Alice",
        "age" => 30,
        "scores" => [95, 88, 92],
        "metadata" => Dict(
            "height" => 155,
            "weight" => 55
        )
    )
    
    # Create a large Dictionary (will use link transport if > 1MB)
    # Generate a larger dataset (~2MB to ensure link transport)
    large_dict = Dict(
        "ids" => collect(1:50000),
        "names" => ["User_$i" for i in 1:50000],
        "scores" => rand(1:100, 50000),
        "categories" => ["Category_$(rand(1:10))" for i in 1:50000],
        "metadata" => Dict(
            "source" => "test_generator",
            "timestamp" => string(Dates.now())
        )
    )
    
    # Test data 1: small Dictionary
    data1 = ("small_dict", small_dict, "dictionary")
    
    # Test data 2: large Dictionary
    data2 = ("large_dict", large_dict, "dictionary")
    
    # Use smartsend with dictionary type
    # For small Dictionary: will use direct transport (JSON encoded)
    # For large Dictionary: will use link transport (uploaded to fileserver)
    env = NATSBridge.smartsend(
        SUBJECT,
        [data1, data2];  # List of (dataname, data, type) tuples
        nats_url = NATS_URL,
        fileserver_url = FILESERVER_URL,
        fileserverUploadHandler = plik_upload_handler,
        size_threshold = 1_000_000,  # 1MB threshold
        correlation_id = correlation_id,
        msg_purpose = "chat",
        sender_name = "dict_sender",
        receiver_name = "",
        receiver_id = "",
        reply_to = "",
        reply_to_msg_id = ""
    )
    
    log_trace("Sent message with $(length(env.payloads)) payloads")
    
    # Log transport type for each payload
    for (i, payload) in enumerate(env.payloads)
        log_trace("Payload $i ('$payload.dataname'):")
        log_trace("  Transport: $(payload.transport)")
        log_trace("  Type: $(payload.type)")
        log_trace("  Size: $(payload.size) bytes")
        log_trace("  Encoding: $(payload.encoding)")
        
        if payload.transport == "link"
            log_trace("  URL: $(payload.data)")
        end
    end
end


# Run the test
println("Starting Dictionary transport test...")
println("Correlation ID: $correlation_id")

# Run sender
println("start smartsend for dictionaries")
test_dict_send()

println("Test completed.")