using NATSBridge
using UUIDs
using JSON
using DataFrames
using Dates

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

# Constants
const NATS_URL = "nats://localhost:4222"
const FILESERVER_URL = "http://localhost:8080"

# Main chat receiver function for scenario 6
function chat_receiver(
    subject::String = "/chat/test";
    nats_url::String = NATS_URL,
    fileserver_url::String = FILESERVER_URL,
    duration::Int = 60,  # Duration in seconds to listen for messages
    max_messages::Int = 100  # Maximum number of messages to receive
)
    println("\n=== Chat Receiver (ServiceB) ===")
    println("Subject: $subject")
    println("NATS URL: $nats_url")
    println("Fileserver URL: $fileserver_url")
    println("Listening duration: $(duration)s")
    println("Max messages: $max_messages")
    println("="^50)
    
    # Create a handler for the fileserver download
    # This will be passed to smartreceive as fileserverDownloadHandler parameter
    fileserverDownloadHandler = (url, max_retries, base_delay, max_delay, correlation_id) -> 
        NATSBridge._fetch_with_backoff(url, max_retries, base_delay, max_delay, correlation_id)
    
    # Connect to NATS and subscribe to the chat subject
    conn = NATS.connect(nats_url)
    
    # Track received messages
    message_count = 0
    total_payloads = 0
    
    # Subscribe to the subject
    subscription = NATS.subscribe(conn, subject)
    
    @info "Listening for chat messages on '$subject'..."
    
    # Listen for messages for the specified duration
    timeout = time() + duration
    
    while time() < timeout && message_count < max_messages
        # Wait for a message with a short timeout
        msg = NATS.waitfor(subscription, 1.0)  # 1 second timeout
        
        if msg !== nothing
            message_count += 1
            println("\n--- Message $(message_count) Received ---")
            
            # Process the message using smartreceive
            payloads = NATSBridge.smartreceive(
                msg;
                fileserverDownloadHandler = fileserverDownloadHandler,
                max_retries = 5,
                base_delay = 100,
                max_delay = 5000
            )
            
            println("Payloads received: $(length(payloads))")
            total_payloads += length(payloads)
            
            # Process each payload
            for (dataname, data, payload_type) in payloads
                println("  - $dataname (type: $payload_type)")
                
                # Handle different types differently for display
                if payload_type == "text"
                    println("    Text content: $(String(data))")
                elseif payload_type == "dictionary"
                    println("    Dictionary content: $(JSON.json(data, 2))")
                elseif payload_type == "table"
                    println("    Table content: $(size(data, 1)) rows, $(size(data, 2)) columns")
                    if size(data, 1) <= 10
                        println("    Sample: $(DataFrames.show(data))")
                    end
                elseif payload_type == "image"
                    println("    Image: $(length(data)) bytes")
                elseif payload_type == "audio"
                    println("    Audio: $(length(data)) bytes")
                elseif payload_type == "video"
                    println("    Video: $(length(data)) bytes")
                elseif payload_type == "binary"
                    println("    Binary: $(length(data)) bytes")
                end
            end
            
            # Extract correlation ID from the message
            json_data = JSON.parse(String(msg.payload))
            println("  Correlation ID: $(json_data["correlationId"])")
            println("  Message ID: $(json_data["msgId"])")
            
            # Optional: Send ACK reply
            reply_to = get(json_data, "replyTo", "")
            if !isempty(reply_to)
                println("  Reply to: $reply_to")
                # Could send an ACK message here
            end
        end
    end
    
    println("\n=== Chat Receiver Summary ===")
    println("Total messages received: $message_count")
    println("Total payloads processed: $total_payloads")
    println("Average payloads per message: $(round(total_payloads / max(message_count, 1), digits=2))")
    println("="^50)
    
    # Cleanup
    NATS.drain(conn)
    
    return message_count
end

# Example usage
if abspath(PROGRAM_FILE) == @__FILE__
    # Parse command line arguments
    if length(ARGS) >= 1
        subject = ARGS[1]
    else
        subject = "/chat/test"
    end
    
    chat_receiver(subject)
end