using JSON, Dates, UUIDs, PrettyPrinting, Base64, DataFrames, NATS, HTTP, Base.Threads
using GeneralUtils, msghandler

# ---------------------------------------------- 100 --------------------------------------------- #
const config = JSON.parsefile("./appconfig.json")
println("Config loaded")

# ----------------------------------------- healthserver ----------------------------------------- #
function healthcheck(req::HTTP.Request)
  return HTTP.Response(200, JSON.json(Dict("status"=>"ok")))
end
const ROUTER = HTTP.Router()
HTTP.register!(ROUTER, "GET", "/ready", healthcheck)
HTTP.register!(ROUTER, "GET", "/health", healthcheck)
HTTP.register!(ROUTER, "GET", "/", healthcheck)
healthserver = HTTP.serve!(ROUTER, "0.0.0.0", 3000)


# ------------------------------------------------------------------------------------------------ #
#                                    registration helper function                                  #
# ------------------------------------------------------------------------------------------------ #
"""
    register_service(connection, service_name, topic, config)

Register a service with the load balancer via NATS request.

# Arguments
- `connection`: NATS connection object
- `service_name`: Name of the service being registered
- `topic`: NATS topic/channel of the registering service
- `config`: Configuration dictionary containing service URLs

# Returns
- Response from the load balancer NATS request

Registers the service with the servicesloadbalancer topic using msghandler.smartpack
to create the registration message envelope.
"""
function register_service(connection, service_name, topic, config)
  register_dict = Dict(
    "register_service_name" => service_name,
    "topic" => topic,
    "free_or_paid" => "free"
  )
  register_payloads = [("register", register_dict, "dictionary")]
  _, register_msg_envelope_json_str = msghandler.smartpack(
              config["externalService"]["servicesloadbalancer"]["url"], 
              register_payloads; 
              msg_purpose="register",   
              broker_url=config["nats_server_info"]["url"],
              fileserver_url=config["externalService"]["fileserver"]["url"])
  try
    reply = NATS.request(connection, 
                      config["externalService"]["servicesloadbalancer"]["url"], 
                      register_msg_envelope_json_str, 
                      timeout=60)
    println("service $service_name is registered")  
  catch
    println("registering service $service_name failed")  
  end
end

ch = Channel(8)
flag = Channel(1)
# Initialize NATS client
connection  = NATS.connect(config["nats_server_info"]["url"])

# use reply() here so that when this_service scale to >1 instance, 
# NATS will randomly select only one receiving instance
sub1 = NATS.subscribe(connection, 
                config["this_service_input_channel"]["nats"][1]; 
                queue_group="group1", spawn=true) do msg
      put!(ch, msg)
  end

# ------------------------------------------------------------------------------------------------ #
#                                      this service main loop                                      #
# ------------------------------------------------------------------------------------------------ #
Threads.@spawn begin  # Register service periodically
  while true
    if isready(flag)
      register_service(connection, 
                      config["this_service_name"],
                      config["this_service_input_channel"]["nats"][1], 
                      config)
    end
    sleep(60) # this line should make spawned thread sleep while allowing main thread running 
  end
end

function main()

# ---------------------------------- register service initially ---------------------------------- #
  register_service(connection, 
                  config["this_service_name"],
                  config["this_service_input_channel"]["nats"][1], 
                  config)
# ---------------------------------------------- 100 --------------------------------------------- #
  println("NATS to HTTP forwarder service for llamacpp_chat_completion service is ready to accept NATS messages at broker=$(config["nats_server_info"]["url"]), topic=$(config["this_service_input_channel"]["nats"][1])")
  while true
    put!(flag, "ready") # ready for new request. do not send register message during llm working
    incoming_nats_msg = take!(ch) # keep this service running.
    println("processing new request")

    # start working, the work may take more than a few minutes. pause async loop registeration to 
    # prevent a loadbalancer assigning new work to this service.
    _ = take!(flag)

    incoming_envelope_json_str = String(incoming_nats_msg.payload)
    incoming_envelope = msghandler.smartunpack(incoming_envelope_json_str)
    if incoming_envelope["msg_purpose"] == "ping"
      outgoing_envelope, outgoing_envelope_json_str = 
      msghandler.smartpack("requester inbox", 
                          [("response", "available", "text")]; 
                          broker_url=config["nats_server_info"]["url"],
                          fileserver_url=config["externalService"]["fileserver"]["url"])
      NATS.publish(connection, incoming_nats_msg.reply_to, outgoing_envelope_json_str)
    else
      incoming_text_array = incoming_envelope["payloads"][1][2]
      body_data = Dict(
          "model" => "embeddinggemma-300M-BF16",
          "input" => incoming_text_array
      )
      headers = [
          "Content-Type" => "application/json",
          "Authorization" => "Bearer no-key"
      ]
      http_response = HTTP.post(config["externalService"]["llamacpp_chat_completion"]["url"], 
                              headers, JSON.json(body_data))

      # HTTP.jl doesn't immediately close connections after post() - it pools them for reuse.
      # Without explicit cleanup, connections accumulate and the llama.cpp server times out
      # on subsequent requests. This forces connection pool cleanup between requests.
      HTTP.close_idle_connections!()
      http_response_body = JSON.parse(http_response.body)
      outgoing_envelope, outgoing_envelope_json_str = 
        msghandler.smartpack("requester inbox", 
                            [("http_response", http_response_body, "dictionary")]; 
                            broker_url=config["nats_server_info"]["url"],
                            fileserver_url=config["externalService"]["fileserver"]["url"])
      # re-register this service to let a loadbalancer knows that this service is ready for next request
      register_service(connection, 
                      config["this_service_name"],
                      config["this_service_input_channel"]["nats"][1], 
                      config)
      NATS.publish(connection, incoming_envelope.reply_to, outgoing_envelope_json_str)
    end
  end
end

main()