-- Purpose: Minimal log watcher using `docker logs -f` as input.
-- Requirements: lua 5.4 + luasocket + docker CLI inside the container.
-- Notes:
-- - Pattern match is plain substring (fast & simple).
-- - Optional Wake-on-LAN is native (no external tools).
-- - Optional port-wait is provided but commented out (mirrors original bash idea).

local socket = require("socket")

local function getenv(name, def)
  local v = os.getenv(name)
  return (v ~= nil and v ~= "") and v or def
end

-- Check if a TCP port is accepting connections within a timeout (seconds)
local function port_is_up(host, port, timeout_sec)
  host = tostring(host or "127.0.0.1")
  port = tonumber(port or 0) or 0
  local timeout = tonumber(timeout_sec or 1) or 1
  if port <= 0 then return false end

  local deadline = socket.gettime() + timeout
  while socket.gettime() < deadline do
    local tcp = socket.tcp()
    if not tcp then return false end
    tcp:settimeout(1)
    local ok = tcp:connect(host, port)
    tcp:close()
    if ok then return true end
    socket.sleep(0.5)
  end
  return false
end

-- ---- Config via env ----
local CONTAINER_NAME 		= getenv("CONTAINER_NAME", "paperless-ai")
local SINCE          		= getenv("SINCE", "0s")
local OLLAMA_HOST    		= getenv("OLLAMA_HOST", "192.168.222.12")
local OLLAMA_PORT    		= tonumber(getenv("OLLAMA_PORT", "11434"))
local SSH_PORT		   		= tonumber(getenv("SSH_PORT", "22"))
local SSH_IDENTITY_FILE = getenv("SSH_IDENTITY_FILE", "/root/.ssh/id_rsa")  -- e.g. "/path/to/id_rsa"
local ERROR_PATTERN  		= getenv(
  "ERROR_PATTERN",
  ("connect EHOSTUNREACH %s:%d"):format(OLLAMA_HOST, OLLAMA_PORT)
)
local FINISH_PATTERN  = getenv("FINISH_PATTERN", "[DEBUG] Finished fetching. Found 0 documents.")  -- e.g. "Server started"

-- Optional Wake-on-LAN
local WOL_MAC   = getenv("WOL_MAC", "")                       -- e.g. "AA:BB:CC:DD:EE:FF"
local WOL_BCAST = getenv("WOL_BCAST", "192.168.222.255")
local WOL_PORT  = tonumber(getenv("WOL_PORT", "9"))

-- Optional: wait for service to come up (kept commented to stay minimal)
-- local UP_WAIT_TIMEOUT = tonumber(getenv("UP_WAIT_TIMEOUT", "90"))

local function log(msg)
  io.stdout:write(os.date("[%F %T] "), msg, "\n"); io.stdout:flush()
end

-- "AA:BB:CC:DD:EE:FF" -> 6 bytes
local function mac_to_bytes(mac)
  local bytes = {}
  for byte in mac:gmatch("(%x%x)") do table.insert(bytes, tonumber(byte, 16)) end
  if #bytes ~= 6 then return nil end
  return string.char(table.unpack(bytes))
end

local function send_wol(mac_str, bcast_ip, port)
  -- Build magic packet
  local bytes = {}
  for byte in mac_str:gmatch("(%x%x)") do table.insert(bytes, tonumber(byte, 16)) end
  if #bytes ~= 6 then return false, "invalid MAC" end
  local mac = string.char(table.unpack(bytes))
  local packet = string.rep(string.char(0xFF), 6) .. mac:rep(16)

  -- Create IPv4 UDP socket (udp4 if available), bind to IPv4 wildcard to lock AF_INET
  local udp = assert((socket.udp4 or socket.udp)())
  udp:settimeout(2)
  assert(udp:setsockname("0.0.0.0", 0))       -- force IPv4 family
  assert(udp:setoption("broadcast", true))    -- allow broadcast

  local ok, err = udp:sendto(packet, bcast_ip, port)
  udp:close()
  return ok ~= nil, err
end

-- Execute a remote command over SSH.
-- Signature must remain: ssh(command, user, host, port, identity_file)
local function ssh(command, user, host, port, identity_file)
  -- Basic validation and defaults
  user = tostring(user or "")
  host = tostring(host or "")
  port = tonumber(port or 22) or 22
  identity_file = tostring(identity_file or "")

  -- Quote a string for safe single-quoted POSIX shell context
  local function sq(s)
    -- Replace ' with: '\''  (close, escape quote, reopen)
    return "'" .. tostring(s):gsub("'", "'\\''") .. "'"
  end

  -- Build base ssh command (run locally)
  -- -oBatchMode to avoid interactive prompts
  -- -oConnectTimeout for faster failure
  -- -oStrictHostKeyChecking uses known_hosts; adjust if needed
  local dest = (user ~= "" and (user .. "@" .. host) or host)
  local pieces = {
    "ssh",
    "-p", tostring(port),
    "-o", "BatchMode=yes",
    "-o", "ConnectTimeout=30",
    "-o", "ServerAliveInterval=5",
    "-o", "ServerAliveCountMax=1",
    "-o", "UserKnownHostsFile=/root/.ssh/known_hosts",
    "-o", "StrictHostKeyChecking=yes",
  }
  if identity_file ~= "" then
    table.insert(pieces, "-i"); table.insert(pieces, identity_file)
  end
  table.insert(pieces, dest)

  -- Quote remote command so the local shell treats it as a single arg.
  table.insert(pieces, '"' .. command .. '"')

  -- Join with spaces for os.execute
  local function join(args)
    -- We only quote the remote command explicitly. Other args are simple tokens.
    return table.concat(args, " ")
  end

  local full = join(pieces)
  log("SSH exec: " .. full)
  local ok, reason, code = os.execute(full)
  if ok == true or ok == 0 then
    log("SSH command completed successfully")
    return true
  else
    local msg = string.format("SSH failed: reason=%s code=%s", tostring(reason), tostring(code))
    log(msg)
    return false, msg
  end
end


local function main()
  log(("Watching container='%s' since='%s'"):format(CONTAINER_NAME, SINCE))
  log(("Looking for pattern: %q"):format(ERROR_PATTERN))

  local cmd = ("docker logs -f --since %q %q 2>&1"):format(SINCE, CONTAINER_NAME)
  local powered_on = false

  while true do

    local fh = assert(io.popen(cmd, "r"))
    for line in fh:lines() do
      -- Plain substring match (no regex)
      if line:find(ERROR_PATTERN, 1, true) ~= nil then
        log(("Detected EHOSTUNREACH for Ollama (%s:%d)."):format(OLLAMA_HOST, OLLAMA_PORT))

        if WOL_MAC ~= "" then
          log(("Sending WOL to %s via %s:%d"):format(WOL_MAC, WOL_BCAST, WOL_PORT))
          local ok, err = send_wol(WOL_MAC, WOL_BCAST, WOL_PORT)
          if ok then
            powered_on = true
            log(("Sucessfully sent WOL to %s via %s:%d"):format(WOL_MAC, WOL_BCAST, WOL_PORT))
          else
            log("WOL failed: " .. tostring(err))
          end
        end

        -- Optional wait (kept commented for minimal parity)
        -- if port_is_up(OLLAMA_HOST, OLLAMA_PORT, UP_WAIT_TIMEOUT) then
        --   log("Ollama reachable again.")
        -- else
        --   log("Timeout waiting for Ollama.")
        -- end

        log("Waiting for SSH to become reachable...")

        if port_is_up(OLLAMA_HOST, SSH_PORT, 60) then
          log("SSH is reachable. Starting ollama service...")
          socket.sleep(2)
          ssh("wsl.exe -d Debian -- 'sudo systemctl enable ollama && sudo systemctl start ollama' && true ", "micro", OLLAMA_HOST, SSH_PORT, SSH_IDENTITY_FILE)
          if (port_is_up(OLLAMA_HOST, OLLAMA_PORT, 90)) then
            log("Ollama service is reachable again.")
            socket.sleep(30)
            break
          else
            log("Timeout waiting for Ollama service to come up after SSH command.")
          end
        end
      end

  if line:find(FINISH_PATTERN, 1, true) ~= nil and powered_on == true then
        log(("Detected finish pattern: %q"):format(FINISH_PATTERN))
        log("Shutting down Ollama host to save power...")
        ssh("wsl.exe -d Debian -- 'sudo systemctl disable ollama && sudo systemctl stop ollama' && true ", "micro", OLLAMA_HOST, SSH_PORT, SSH_IDENTITY_FILE)
        ssh("shutdown.exe /s /t 0", "micro", OLLAMA_HOST, SSH_PORT, SSH_IDENTITY_FILE)
        socket.sleep(5)
        powered_on = false
        break
      end
    end

    fh:close()
    log("Log stream ended.")
  end

end

main()