import Control.Concurrent (forkIO)
import Control.Exception (catch, try, SomeException)

import Control.Concurrent (forkIO, MVar, newEmptyMVar, putMVar, takeMVar, killThread, forkFinally, isEmptyMVar)
import Control.Exception (catch, try, SomeException, IOException)

import Control.Monad (forM_, unless, forever)

import Control.Monad (forM_, unless, forever, void, when)

import System.Posix.Signals

import System.IO.Error (isDoesNotExistError)

  stopFlag <- newEmptyMVar -- Create a flag that the signal handler can set
  -- Install signal handlers for SIGTERM (systemd stop) and SIGINT (Ctrl+C)
  installHandler sigTERM (Catch $ putMVar stopFlag ()) Nothing

              -- Clean up old socket if exists
          catch (removeFile socketPath) (\(_::SomeException) -> pure ())

          -- Clean up any existing socket file
          catch (removeFile socketPath)
                (\(e :: IOException) ->
                  if isDoesNotExistError e then pure () else ioError e)
          -- Create and bind the UNIX domain socket

          listen sock 1

          listen sock 10
            -- Run the accept loop in a separate thread
          acceptThread <- forkIO $
            acceptLoop sock stopFlag (\conn -> handleClient conn ip raw)
          -- Wait until SIGTERM
          takeMVar stopFlag
          putStrLn "Received SIGTERM, shutting down gracefully..."
          --clear the iptables
          stop ip raw

          forever $ do
            (conn, _) <- accept sock
            _ <- forkIO $ (handleClient conn ip raw)
            pure ()

          -- Stop accepting, clean up
          killThread acceptThread
          close sock
          removeFile socketPath

-- | Loop accepting connections until stopFlag is set
acceptLoop :: Socket -> MVar () -> (Socket -> IO ()) -> IO ()
acceptLoop sock stopFlag handleConn = do
  let loop = do
        stillRunning <- isEmptyMVar stopFlag
        when stillRunning $ do
          (conn, _) <- accept sock
          void $ forkFinally (handleConn conn) (\_ -> close conn)
          loop
  loop