-- KumoMTA configuration
local kumo = require 'kumo'
local utils = require 'policy-extras.policy_utils'

-- Add regex set for PMTA-style backoff pattern matching
local regex_set = kumo.regex_set_map.new()
local common_errors_regex = kumo.regex_set_map.new()
local blocking_errors_regex = kumo.regex_set_map.new()

-- Load authentication users directly in the code for simplicity
local smtp_auth_users = {
  ['user1'] = 'password1',
  ['user2'] = 'password2',
  -- Add more users as needed
}

-- Load the policy helpers
local shaping = require 'policy-extras.shaping'
local queue_module = require 'policy-extras.queue'
local sources = require 'policy-extras.sources'
local dkim_sign = require 'policy-extras.dkim_sign'
local rollup = require 'policy-extras.rollup'

-- Configure sources and pools
sources:setup { '/opt/kumomta/etc/policy/sources.toml' }

-- Configure traffic shaping with automation
local shaper = shaping:setup_with_automation {
  publish = { 'http://127.0.0.1:8008' },
  subscribe = { 'http://127.0.0.1:8008' },
  extra_files = { 
    '/opt/kumomta/share/policy-extras/shaping.toml',
    '/opt/kumomta/share/community/shaping.toml',
    '/opt/kumomta/etc/policy/shaping.toml'
  },
}

-- Configure queue management
local queue_helper = queue_module:setup { '/opt/kumomta/etc/policy/queues.toml' }

-- Configure DKIM signing
local dkim_signer = dkim_sign:setup { '/opt/kumomta/etc/policy/dkim_data.toml' }

-- Track domains that are in backoff mode with timestamps
local domains_in_backoff = {}

-- Common errors from your PMTA config
local common_error_patterns = {
  ["generating high volumes of.* "] = true,
  ["Excessive unknown recipients - possible Open Relay"] = true,
  ["^421 .* too many errors"] = true,
  ["blocked.*spamhaus"] = true,
  ["451 Rejected"] = true
}

-- Blocking errors (more serious)
local blocking_error_patterns = {
  ["421 .* SERVICE NOT AVAILABLE"] = true,
  ["generating high volumes of.* complaints from AOL"] = true,
  ["554 .*aol.com"] = true,
  ["421dynt1"] = true,
  ["HVU:B1"] = true,
  ["DNS:NR"] = true,
  ["RLY:NW"] = true,
  ["DYN:T1"] = true,
  ["RLY:BD"] = true,
  ["RLY:CH2"] = true,
  ["421 .* Please try again later"] = true,
  ["421 Message temporarily deferred"] = true,
  ["VS3-IP5 Excessive unknown recipients"] = true,
  ["VSS-IP Excessive unknown recipients"] = true,
  ["\\[GL01\\] Message from"] = true,
  ["\\[TS01\\] Messages from"] = true,
  ["\\[TS02\\] Messages from"] = true,
  ["\\[TS03\\] All messages from"] = true,
  ["exceeded the rate limit"] = true,
  ["exceeded the connection limit"] = true,
  ["Mail rejected by Windows Live Hotmail for policy reasons"] = true,
  ["mail.live.com\\/mail\\/troubleshooting.aspx"] = true,
  ["421 Message Rejected"] = true,
  ["Client host rejected"] = true,
  ["blocked using UCEProtect"] = true,
  ["Mail Refused"] = true,
  ["421 Exceeded allowable connection time"] = true,
  ["amIBlockedByRR"] = true,
  ["block-lookup"] = true,
  ["Too many concurrent connections from source IP"] = true,
  
  -- Additional blocking errors
  ["too many"] = true,
  ["Exceeded allowable connection time"] = true,
  ["Connection rate limit exceeded"] = true,
  ["refused your connection"] = true,
  ["try again later"] = true,
  ["try later"] = true,
  ["550 RBL"] = true,
  ["TDC internal RBL"] = true,
  ["connection refused"] = true,
  ["please see www.spamhaus.org"] = true,
  ["Message Rejected"] = true,
  ["Delivery report"] = true,
  ["refused by antispam"] = true,
  ["Service not available"] = true,
  ["currently blocked"] = true,
  ["locally blacklisted"] = true,
  ["not currently accepting mail from your ip"] = true,
  ["421.*closing connection"] = true,
  ["421.*Lost connection"] = true,
  ["476 connections from your host are denied"] = true,
  ["421 Connection cannot be established"] = true,
  ["421 temporary envelope failure"] = true,
  ["421 4.4.2 Timeout while waiting for command"] = true,
  ["450 Requested action aborted"] = true,
  ["550 Access denied"] = true,
  ["exceeded the rate limit"] = true,
  ["421rlynw"] = true,
  ["permanently deferred"] = true,
  ["\\d+\\.\\d+\\.\\d+\\.\\d+ blocked"] = true,
  ["www\\.spamcop\\.net\\/bl\\.shtml"] = true,
  
  -- Original SMTPRESPONS patterns
  ["421 PR\\(ct1\\)"] = true,
  ["^550 SC-001"] = true,
  ["420 Resources unavailable temporarily"] = true,
  ["^Resources unavailable temporarily"] = true,
  ["^421"] = true,
  ["^450"] = true,
  ["^try later"] = true,
  ["^553"] = true,
  ["^554"] = true,
  ["^busy"] = true,
  ["^WSAECONNREFUSED"] = true,
  ["^WSAECONNRESET"] = true,
  ["^Connection attempt failed"] = true
}



local regex_patterns = {
  -- AOL errors
  ["421 .* SERVICE NOT AVAILABLE"] = true,
  ["generating high volumes of.* complaints from AOL"] = true,
  ["554 .*aol.com"] = true,
  ["421dynt1"] = true,
  ["HVU:B1"] = true,
  ["DNS:NR"] = true,
  ["RLY:NW"] = true,
  ["DYN:T1"] = true,
  ["RLY:BD"] = true,
  ["RLY:CH2"] = true,
  
  -- Yahoo errors
  ["421 .* Please try again later"] = true,
  ["421 Message temporarily deferred"] = true,
  ["VS3-IP5 Excessive unknown recipients"] = true,
  ["VSS-IP Excessive unknown recipients"] = true,
  ["\\[GL01\\] Message from"] = true,
  ["\\[TS01\\] Messages from"] = true,
  ["\\[TS02\\] Messages from"] = true,
  ["\\[TS03\\] All messages from"] = true,
  
  -- Hotmail errors
  ["exceeded the rate limit"] = true,
  ["exceeded the connection limit"] = true,
  ["Mail rejected by Windows Live Hotmail for policy reasons"] = true,
  ["mail.live.com\\/mail\\/troubleshooting.aspx"] = true,
  
  -- Adelphia errors
  ["421 Message Rejected"] = true,
  ["Client host rejected"] = true,
  ["blocked using UCEProtect"] = true,
  
  -- Road Runner errors
  ["Mail Refused"] = true,
  ["421 Exceeded allowable connection time"] = true,
  ["amIBlockedByRR"] = true,
  ["block-lookup"] = true,
  ["Too many concurrent connections from source IP"] = true,
  
  -- General errors
  ["too many"] = true,
  ["Exceeded allowable connection time"] = true,
  ["Connection rate limit exceeded"] = true,
  ["refused your connection"] = true,
  ["try again later"] = true,
  ["try later"] = true,
  ["550 RBL"] = true,
  ["TDC internal RBL"] = true,
  ["connection refused"] = true,
  ["please see www.spamhaus.org"] = true,
  ["Message Rejected"] = true,
  ["Delivery report"] = true,
  ["refused by antispam"] = true,
  ["Service not available"] = true,
  ["currently blocked"] = true,
  ["locally blacklisted"] = true,
  ["not currently accepting mail from your ip"] = true,
  ["421.*closing connection"] = true,
  ["421.*Lost connection"] = true,
  ["476 connections from your host are denied"] = true,
  ["421 Connection cannot be established"] = true,
  ["421 temporary envelope failure"] = true,
  ["421 4.4.2 Timeout while waiting for command"] = true,
  ["450 Requested action aborted"] = true,
  ["550 Access denied"] = true,
  ["exceeded the rate limit"] = true,
  ["421rlynw"] = true,
  ["permanently deferred"] = true,
  ["\\d+\\.\\d+\\.\\d+\\.\\d+ blocked"] = true,
  ["www\\.spamcop\\.net\\/bl\\.shtml"] = true,
  ["generating high volumes of.*"] = true,
  ["Excessive unknown recipients - possible Open Relay"] = true,
  ["^421 .* too many errors"] = true,
  ["blocked.*spamhaus"] = true,
  ["451 Rejected"] = true,
  
  -- Original SMTPRESPONS patterns
  ["421 PR\\(ct1\\)"] = true,
  ["^550 SC-001"] = true,
  ["420 Resources unavailable temporarily"] = true,
  ["^Resources unavailable temporarily"] = true,
  ["^421"] = true,
  ["^450"] = true,
  ["^try later"] = true,
  ["^553"] = true,
  ["^554"] = true,
  ["^busy"] = true,
  ["^WSAECONNREFUSED"] = true,
  ["^WSAECONNRESET"] = true,
  ["^Connection attempt failed"] = true
}

-- Now create the regex sets with the prepared patterns
local regex_set = kumo.regex_set_map.new(regex_patterns)
local common_errors_regex = kumo.regex_set_map.new(common_error_patterns)
local blocking_errors_regex = kumo.regex_set_map.new(blocking_error_patterns)

-- Create tracking tables for backoff
local domains_in_common_backoff = {}
local domains_in_blocking_backoff = {}
local domain_bounce_handlers = {}
local auth_failures = {}


-- Auth attempt rate limiting
local function check_auth_rate_limit(ip_addr)
  local now = os.time()
  local failures = auth_failures[ip_addr] or {}
  
  -- Clean up old entries
  local recent_failures = 0
  for i = #failures, 1, -1 do
      if now - failures[i] > 600 then  -- 10 minute window
          table.remove(failures, i)
      else
          recent_failures = recent_failures + 1
      end
  end
  
  -- Check if rate limit exceeded (10 failures in 10 minutes)
  if recent_failures >= 10 then
      return true
  end
  
  return false
end

-- Record auth failure
local function record_auth_failure(ip_addr)
  local now = os.time()
  auth_failures[ip_addr] = auth_failures[ip_addr] or {}
  table.insert(auth_failures[ip_addr], now)
end

-- Default bounce handler for all domains
local function default_bounce_handler(msg, bounce_category, destination_info)
  local sender = msg:sender().email
  local recipient = msg:recipient().email
  local domain = msg:recipient().domain
  local site_name = destination_info.site_name
  
  kumo.log_debug("Processing bounce - Category:", bounce_category, 
                "Sender:", sender, 
                "Recipient:", recipient,
                "Domain:", domain,
                "Site name:", site_name)
  
  -- Classify bounce types for handling
  local is_hard_bounce = false
  local is_reputation_issue = false
  local is_temp_failure = false
  
  -- Hard bounces
  if bounce_category == "BadMailbox" or 
     bounce_category == "InvalidRecipient" or 
     bounce_category == "BadDomain" then
    is_hard_bounce = true
  
  -- Reputation issues
  elseif bounce_category == "SpamRelated" or 
         bounce_category == "PolicyRelated" or 
         bounce_category == "IPReputation" or
         bounce_category == "RateLimited" then
    is_reputation_issue = true
  
  -- Temporary failures
  elseif bounce_category == "QuotaIssues" or 
         bounce_category == "DeferredTemporarily" or
         bounce_category == "MessageExpired" then
    is_temp_failure = true
  end
  
  -- Log according to type
  if is_hard_bounce then
    kumo.log_info("Hard bounce detected for:", recipient, "Category:", bounce_category)
  elseif is_reputation_issue then
    kumo.log_info("Reputation issue detected for domain:", domain, "Category:", bounce_category)
  elseif is_temp_failure then
    kumo.log_info("Temporary failure for:", recipient, "Category:", bounce_category)
  else
    kumo.log_info("General bounce for:", recipient, "Category:", bounce_category)
  end
  
  -- Default response is to let normal processing continue
  return false
end

-- Register a domain-specific bounce handler
function register_domain_bounce_handler(domain, handler_function)
  domain_bounce_handlers[domain] = handler_function
  kumo.log_debug("Registered custom bounce handler for domain:", domain)
end

-- Process bounce with appropriate handler
function process_bounce(msg, bounce_category, destination_info)
  local domain = msg:recipient().domain
  
  -- Check if we have a specific handler for this domain
  if domain_bounce_handlers[domain] then
    kumo.log_debug("Using custom bounce handler for domain:", domain)
    return domain_bounce_handlers[domain](msg, bounce_category, destination_info)
  end
  
  -- Fall back to default handler
  return default_bounce_handler(msg, bounce_category, destination_info)
end


local function handle_bounce_by_category(msg, bounce_category)
  local sender = msg:sender().email
  local recipient = msg:recipient().email
  local domain = msg:recipient().domain
  
  -- Log bounce information with category
  kumo.log_info("Bounce detected - Category:", bounce_category, 
                "Sender:", sender, 
                "Recipient:", recipient,
                "Domain:", domain)
  
  -- You can implement different actions based on bounce category
  if bounce_category == "BadMailbox" or bounce_category == "InvalidRecipient" then
    -- These are hard bounces - you might want to add to suppression list
    kumo.log_info("Hard bounce detected, should update suppression list for:", recipient)
    
    -- Here you could call an API or update a database
    -- Example pseudocode:
    -- update_suppression_list(recipient, bounce_category)
  elseif bounce_category == "SpamRelated" or bounce_category == "PolicyRelated" then
    -- These may indicate sending reputation issues
    kumo.log_info("Reputation-related bounce for domain:", domain)
    
    -- Consider implementing domain throttling or backoff
    -- Example pseudocode:
    -- apply_domain_throttling(domain)
  elseif bounce_category == "QuotaIssues" or bounce_category == "MessageExpired" then
    -- These are soft bounces - temporary issues
    kumo.log_info("Soft bounce detected for:", recipient)
    
    -- You might want to track these for repeated soft bounces
    -- Example pseudocode:
    -- increment_soft_bounce_counter(recipient)
  end
  
  return false -- Let normal processing continue
end

-- SMTP Authentication handler with proper logging
kumo.on('smtp_server_auth_plain', function(authz, authc, password, conn_meta)
  local peer_ip = conn_meta:get_meta('peer_address')
  local is_using_tls = conn_meta:get_meta('tls_status') == 'Active'
  
  kumo.log_debug("Auth attempt - IP:", peer_ip, 
                  "User:", authc or "<empty>", 
                  "TLS:", is_using_tls and "Yes" or "No")
  kumo.log_debug("Auth attempt from:", conn_meta:get_meta('peer_address'))
  kumo.log_debug("Auth user:", authc or "nil")
  kumo.log_debug("Password length:", string.len(password or ""))

  -- Enforce TLS for authentication
  if not is_using_tls then
    kumo.log_warn("Auth attempt without TLS from:", peer_ip)
    -- Sleep to slow down brute force attempts
    kumo.sleep(3)
    return false
  end
  
  -- Check rate limiting
  if check_auth_rate_limit(peer_ip) then
    kumo.log_warn("Auth rate limit exceeded for IP:", peer_ip)
    -- Sleep to slow down brute force attempts
    kumo.sleep(3)
    return false
  end


  -- Validate credentials
  if not authc or authc == "" then
      kumo.log_debug("Auth failed: empty username from IP:", peer_ip)
      record_auth_failure(peer_ip)
      return false
  end
  
  if not password or password == "" then
      kumo.log_debug("Auth failed: empty password from IP:", peer_ip)
      record_auth_failure(peer_ip)
      return false
  end
  
  local is_authorized = smtp_auth_users[authc] == password
  
  if is_authorized then
    kumo.log_info("Auth successful for user:", authc, "from IP:", peer_ip)
    -- Set metadata for successful auth
    conn_meta:set_meta('authz_id', authc)
    conn_meta:set_meta('authenticated', true)
    conn_meta:set_meta('auth_time', os.time())
  else
      kumo.log_warn("Auth failed for user:", authc, "from IP:", peer_ip)
      record_auth_failure(peer_ip)
      -- Sleep to slow down brute force attempts
      kumo.sleep(1)
  end
  
  return is_authorized
end)

-- HTTP Basic Auth validator
kumo.on('http_server_validate_auth_basic', function(user, pass)
  return smtp_auth_users[user] == pass
end)

-- Handler for SMTP client responses to trigger backoff behavior
kumo.on('smtp_client_response_received', function(response, destination, msg)
  local domain = destination.domain
  local site_name = destination.site_name
  local tenant = msg:get_meta('tenant')
  local current_time = os.time()

  local bounce_category = msg:get_meta('bounce_classification')
  local response_code = tonumber(string.match(response, "^(%d%d%d)") or "0")

  -- Only process as bounce if it's a 4xx or 5xx response
  if response_code >= 400 then
    -- Get or compute bounce classification
    local bounce_category = msg:get_meta('bounce_classification') or "Uncategorized"
    
    -- Create destination info object with additional context
    local destination_info = {
      domain = destination.domain,
      site_name = destination.site_name,
      mx_hostname = destination.mx_hostname,
      response = response,
      response_code = response_code,
    }
    
    -- Process the bounce with our framework
    local handled = process_bounce(msg, bounce_category, destination_info)
    if handled then
      return true -- Skip further processing if the handler took care of it
    end
  end

  if bounce_category and bounce_category ~= "Uncategorized" then
    return handle_bounce_by_category(msg, bounce_category)
  end
  
  -- Check for common errors (shorter backoff duration)
  if common_errors_regex[response] then
      kumo.log_debug("Common error backoff triggered for domain:", domain, "site:", site_name, "response:", response)
      
      -- Record domain in common backoff state
      domains_in_common_backoff[domain] = current_time
      
      -- Apply moderate retry delay
      msg:set_meta("retry_interval", "15 minutes")
      
      return true
  
  -- Check for blocking errors (longer backoff duration)
  elseif blocking_errors_regex[response] then
      kumo.log_debug("Blocking error backoff triggered for domain:", domain, "site:", site_name, "response:", response)
      
      -- Record domain in blocking backoff state
      domains_in_blocking_backoff[domain] = current_time
      
      -- Apply stronger retry delay
      msg:set_meta("retry_interval", "30 minutes")
      
      return true
  end
  
  return false -- Let normal processing continue
end)

-- Called On Startup, handles initial configuration
kumo.on('init', function()
  -- Set diagnostic log level to see our log messages
  kumo.set_diagnostic_log_filter 'kumod=info,smtp_server=debug,auth=debug,lua=debug,tls=info'
  
  -- Define spools for message storage
  kumo.define_spool {
    name = 'data',
    path = '/var/spool/kumomta/data',
    kind = 'RocksDB',
  }

  kumo.define_spool {
    name = 'meta',
    path = '/var/spool/kumomta/meta',
    kind = 'RocksDB',
  }

  -- Configure publishing of logs to automation daemon
  shaper.setup_publish()

  kumo.spawn_task {
    event_name = 'auth_cleanup_task',
    args = {},
  }

  kumo.configure_local_logs {
    log_dir = '/var/log/kumomta',
    max_segment_duration = '60 minutes',  -- Create new log files every hour
    -- max_file_size = 1024 * 1024 * 1024,                -- Or when they reach 1GB uncompressed
    
    -- Log additional headers and metadata
    headers = { 
      'Subject', 
      'Message-ID', 
      'X-Campaign-ID',
      'From',
      'To',
      'Sender'
    },
    
    -- Log specific metadata fields
    meta = { 
      'tenant', 
      'campaign', 
      'authenticated_sender',
      'auth_type',
      'source_ip'
    },
    
    -- Configure per-record settings
    per_record = {
      
      
      -- Special settings for delivery logs
      Delivery = {
        log_dir = '/var/log/kumomta/delivery',  -- Store delivery logs separately
      },
      
      -- Special settings for reception logs
      Reception = {
        suffix = '_reception',  -- Add suffix to filenames
      },
      
      -- Special settings for bounce logs
      Bounce = {
        log_dir = '/var/log/kumomta/bounces',  
      },
      
      -- Special settings for feedback reports
      Feedback = {
        log_dir = '/var/log/kumomta/feedback',
      },
    },
  }

  kumo.configure_log_hook {
    name = 'critical_events',
    
    -- Only log these event types
    per_record = {
      -- Only these record types will be hooked
      Feedback = { enable = true },
      Bounce = { enable = true },
      
    },
  }

  -- Register a handler for gmail.com
  register_domain_bounce_handler("gmail.com", function(msg, category, dest_info)
    -- Gmail-specific bounce processing
    if category == "RateLimited" then
      kumo.log_info("Gmail rate limit detected, applying aggressive backoff")
      -- Apply stronger backoff for Gmail rate limits
      return true -- Indicate we've handled this bounce
    end
    return false -- Let default handler process it
  end)
  
  -- Register a handler for yahoo.com
  register_domain_bounce_handler("yahoo.com", function(msg, category, dest_info)
    -- Yahoo-specific bounce processing
    if category == "SpamRelated" then
      kumo.log_info("Yahoo spam complaint detected, applying throttling")
      -- Apply specific throttling for Yahoo spam complaints
      return true
    end
    return false
  end)

  -- Configure bounce classification
  kumo.configure_bounce_classifier {
    files = {
      '/opt/kumomta/share/bounce_classifier/iana.toml',
      '/opt/kumomta/etc/policy/custom_bounce_classifier.toml',
    },
    -- Add these improved settings
    pool_size = math.floor(kumo.available_parallelism() / 2),  -- Use half of available cores
    cache_size = 2048,  -- Larger cache for better performance
    uncategorized_cache_size = 1024, -- Cache for unclassified results
  }

  -- Configure HTTP listener for admin access
  kumo.start_http_listener {
    listen = '0.0.0.0:8000',
    trusted_hosts = { '127.0.0.1', '::1' },
  }

  -- Configure SMTP listener for port 25
  kumo.start_esmtp_listener {
    listen = '0.0.0.0:25',
    hostname = 'kumo.s-gii.com',
    relay_hosts = { '127.0.0.1'},  -- ONLY localhost can relay without auth
    tls_certificate = '/etc/letsencrypt/live/kumo.s-gii.com/fullchain.pem',
    tls_private_key = '/etc/letsencrypt/live/kumo.s-gii.com/privkey.pem',
    trace_headers = {
      received_header = true,
      supplemental_header = true,
      header_name = 'X-KumoRef',
      include_meta_names = { 'tenant', 'campaign' }, -- Add metadata you want to track
    },
    -- Customized settings for submission port
    banner = 'KumoMTA Secure SMTP Submission Service Ready',
    client_timeout = '5 minutes',                 -- Longer timeout for submission
    max_message_size = 50 * 1024 * 1024,          -- 50MB max message size
    max_recipients_per_message = 500,             -- More recipients allowed
    max_messages_per_connection = 1000,           -- More messages per connection
    invalid_line_endings = 'Fix',     
  }
  
  -- Add submission port (587) for authenticated sending
  kumo.start_esmtp_listener {
    listen = '0.0.0.0:587',
    hostname = 'kumo.s-gii.com',
    relay_hosts = { '127.0.0.1' },  -- ONLY localhost can relay without auth
    tls_certificate = '/etc/letsencrypt/live/kumo.s-gii.com/fullchain.pem',
    tls_private_key = '/etc/letsencrypt/live/kumo.s-gii.com/privkey.pem',
    trace_headers = {
      received_header = true,
      supplemental_header = true,
      header_name = 'X-KumoRef',
      include_meta_names = { 'tenant', 'campaign' },
    },
  }
  
  -- Spawn a task to clean up old backoff entries
  kumo.spawn_task {
    event_name = 'clean_backoff_entries',
    args = {},
  }

  
end)

-- Add this event handler if using log hooks
kumo.on('should_enqueue_log_record', function(msg, record_json, hook_name)
  if hook_name == 'critical_events' then
    -- Queue to a specific domain for processing by your webhook handler
    msg:set_meta('queue', 'critical_logs.internal')
    return true
  end
  return false
end)


-- Clean up old backoff entries periodically
kumo.on('clean_backoff_entries', function()
  while true do
      local current_time = os.time()
      local common_domains_to_remove = {}
      local blocking_domains_to_remove = {}
      
      -- Process common error backoffs (expire after 30 minutes)
      for domain, backoff_time in pairs(domains_in_common_backoff) do
          if (current_time - backoff_time) > 1800 then
              table.insert(common_domains_to_remove, domain)
          end
      end
      
      -- Process blocking error backoffs (expire after 2 hours)
      for domain, backoff_time in pairs(domains_in_blocking_backoff) do
          if (current_time - backoff_time) > 7200 then
              table.insert(blocking_domains_to_remove, domain)
          end
      end
      
      -- Remove expired domains
      for _, domain in ipairs(common_domains_to_remove) do
          kumo.log_debug("Removing domain from common error backoff:", domain)
          domains_in_common_backoff[domain] = nil
      end
      
      for _, domain in ipairs(blocking_domains_to_remove) do
          kumo.log_debug("Removing domain from blocking error backoff:", domain)
          domains_in_blocking_backoff[domain] = nil
      end
      
      -- Check every 5 minutes
      kumo.sleep(300)
  end
end)

-- Configure traffic shaping
kumo.on('get_egress_path_config', shaper.get_egress_path_config)

-- Extend queue config to apply backoff settings
kumo.on('get_queue_config', function(domain, tenant, campaign, routing_domain)
  local params = {}
  local current_time = os.time()
  

  queue_helper:apply_to_params(domain, tenant, campaign, routing_domain, params)

  if domain == 'critical_logs.internal' then
    return kumo.make_queue_config {
      protocol = {
        custom_lua = {
          constructor = 'make.webhook',
        },
      },
    }
  end

  -- Check if domain is in common error backoff (shorter duration, gentler settings)
  if domains_in_common_backoff[domain] then
      local backoff_time = domains_in_common_backoff[domain]
      
      -- If backoff was triggered within the last 30 minutes
      if (current_time - backoff_time) < 1800 then
          kumo.log_debug("Applying common error backoff settings for domain:", domain)
          params.retry_interval = "15 minutes"
          params.max_retry_interval = "30 minutes"
          params.max_message_rate = "20/minute"
      else
          -- Clear old backoff entries
          domains_in_common_backoff[domain] = nil
      end
  end
  
  -- Check if domain is in blocking error backoff (longer duration, stricter settings)
  if domains_in_blocking_backoff[domain] then
      local backoff_time = domains_in_blocking_backoff[domain]
      
      -- If backoff was triggered within the last 2 hours
      if (current_time - backoff_time) < 7200 then
          kumo.log_debug("Applying blocking error backoff settings for domain:", domain)
          params.retry_interval = "30 minutes"
          params.max_retry_interval = "60 minutes"
          params.max_message_rate = "5/minute"
          
          -- More restrictive settings than common errors
          params.max_connection_rate = "1/minute"
      else
          -- Clear old backoff entries
          domains_in_blocking_backoff[domain] = nil
      end
    end

    rollup.apply_ip_rollup_to_queue_config(domain, routing_domain, params)
  
  return kumo.make_queue_config(params)
end)

-- Handle domain permissions for relay
kumo.on('get_listener_domain', function(domain, listener, conn_meta)
  local authz_id = conn_meta:get_meta('authz_id')
  local peer_ip = conn_meta:get_meta('peer_address')
  local is_using_tls = conn_meta:get_meta('tls_status') == 'Active'
  
  kumo.log_debug("Checking relay permissions - Domain:", domain,
                  "IP:", peer_ip,
                  "Auth:", authz_id or "none",
                  "TLS:", is_using_tls and "Yes" or "No")
  kumo.log_debug("Connection authenticated as:", authz_id or "not authenticated")
  
  -- Allow authenticated users to relay to any domain
  if authz_id then
    kumo.log_debug("Allowing relay for authenticated user:", authz_id)
    return kumo.make_listener_domain {
      relay_to = true,
    }
  end

  -- Special handling for feedback loop domain
  if domain == "fbl.kumo.s-gii.com" then
    kumo.log_debug("FBL domain detected:", domain)
    return kumo.make_listener_domain {
      log_arf = "LogThenDrop",  -- Process and log ARF messages, then drop them
      relay_to = false,
    }
  end
  
  -- Add other feedback loop domains as needed
  if domain == "abuse.kumo.s-gii.com" then
    return kumo.make_listener_domain {
      log_arf = "LogThenDrop",
      relay_to = false,
    }
  end
  
  -- For localhost connections, allow relay to specific domains
  if peer_ip == "127.0.0.1" or peer_ip == "::1" then
    -- Allow relay to internal domains
    if domain == "kumo.s-gii.com" then
        return kumo.make_listener_domain {
            relay_to = true,
        }
    end
  end

  kumo.log_debug("Denying relay permission for unauthenticated connection to:", domain)
  return nil
end)

-- Processing of incoming messages via SMTP
kumo.on('smtp_server_message_received', function(msg, conn_meta)
  -- Protect against SMTP Smuggling
  local failed = msg:check_fix_conformance(
    'NON_CANONICAL_LINE_ENDINGS',
    ''
  )
  if failed then
    kumo.reject(552, string.format('5.6.0 %s', failed))
  end
  
  -- If we're authenticated, set the tenant based on the authenticated user
  local authz_id = conn_meta:get_meta('authz_id')
  local peer_ip = conn_meta:get_meta('peer_address')
  local is_authenticated = conn_meta:get_meta('authenticated') or false
  local is_using_tls = conn_meta:get_meta('tls_status') == 'Active'

  -- Set metadata based on authentication
  if authz_id then
    kumo.log_debug("Processing message from authenticated user:", authz_id)
    msg:set_meta('tenant', authz_id)
    msg:set_meta('authenticated_sender', authz_id)
    msg:set_meta('auth_type', 'plain')
  end

  -- Import useful headers for logging
  msg:import_x_headers {
    'subject',
    'message-id',
    'x-campaign-id',
    'x-tenant'
  }
  
  -- Call the queue helper to set up the queue for the message
  queue_helper:apply(msg)
  
  -- Set source IP as metadata for tracking
  msg:set_meta('source_ip', peer_ip)
  msg:set_meta('using_tls', is_using_tls and "yes" or "no")

  -- DKIM signing should come last
  dkim_signer(msg)
end)

-- Same processing for HTTP-generated messages
kumo.on('http_message_generated', function(msg)
  -- Call the queue helper to set up the queue for the message
  queue_helper:apply(msg)
  
  -- DKIM signing should come last
  dkim_signer(msg)
end)

-- Implement the webhook handler
kumo.on('make.webhook', function(domain, tenant, campaign)
  local client = kumo.http.build_client {}
  
  local sender = {}
  function sender:send(message)
    local request = client:post 'https://kumo.s-gii.com/logs'
    
    request:header('Content-Type', 'application/json')
    request:body(message:get_data())
    
    local response = request:send()
    if response:status_is_success() then
      return string.format('250 Webhook delivered: %s', response:status_reason())
    end
    
    kumo.reject(450, string.format('4.0.0 Webhook failed: %s', response:text()))
  end
  
  function sender:close()
    -- Nothing to do
  end
  
  return sender
end)

-- Task implementation
kumo.on('auth_cleanup_task', function()
  while true do
      kumo.log_debug("Running auth failure cleanup task")
      local now = os.time()
      local removed = 0
      
      -- Clean up auth failures older than 1 hour
      for ip, failures in pairs(auth_failures) do
          local new_failures = {}
          for _, time in ipairs(failures) do
              if now - time < 3600 then  -- 1 hour
                  table.insert(new_failures, time)
              else
                  removed = removed + 1
              end
          end
          
          if #new_failures == 0 then
              auth_failures[ip] = nil
          else
              auth_failures[ip] = new_failures
          end
      end
      
      kumo.log_debug("Auth cleanup: removed", removed, "entries")
      
      -- Run every 10 minutes
      kumo.sleep(600)
  end
end)