module Main (main) where

import ConfigParser
import ConfigParserUtils (getYamlConfigEntry)
import Env
import Tproxy
import Ip
import Udp2Raw
import Network.Socket
import Control.Concurrent
    ( forkIO, MVar, newEmptyMVar, putMVar, takeMVar, killThread, isEmptyMVar )
import Control.Exception (catch, IOException)
import System.Environment (getArgs)
import System.IO
    ( hPutStrLn, stderr, hSetBuffering, BufferMode(LineBuffering, NoBuffering)
    , hGetLine, hClose, stdout, IOMode(ReadWriteMode), Handle )
import System.Posix.Files (setFileMode)
import System.Posix.Signals
import System.Directory (removeFile)
import Data.Maybe (fromMaybe, listToMaybe)
import Data.List (find)

import Control.Concurrent
import Control.Concurrent.Async
import Control.Monad (forever)

import qualified Data.Text as T

socketPath :: FilePath
socketPath = "/run/tproxy/socket"

main :: IO ()
main = do
    hSetBuffering stdout NoBuffering

    stopFlag <- newEmptyMVar
    installHandler sigTERM (Catch $ putMVar stopFlag ()) Nothing

    -- Get the server list from command line
    args <- getArgs
    let configFile = case args of
                        (f:_) -> f
                        []    -> "servers.yaml"

    -- Try to load .env first
    let environment = "/run/tproxy/.env"
    eEnv <- Env.loadEnv environment
    case eEnv of
        Right envVal -> do
            putStrLn $ "Loaded environment: " ++ show envVal
            start (Env.ip envVal) (Env.rawMode envVal)
        Left _ -> do
            -- Fallback to YAML
            yamlCfg <- parseConfigFile configFile
            case yamlCfg of
                Right (entry:_) -> do
                    ip <- startEntry entry environment
                    putStrLn ("started on " ++ ip)
                Right [] -> putStrLn "YAML configuration is empty"
                Left err -> hPutStrLn stderr $ "Failed to load .env and YAML: " ++ err

    -- Clean up stale socket
    catch (removeFile socketPath) (\(_ :: IOException) -> pure ())
    sock <- socket AF_UNIX Stream 0
    bind sock (SockAddrUnix socketPath)
    setFileMode socketPath 0o775
    listen sock 10

    acceptThread <- forkIO $ acceptLoop sock stopFlag configFile environment

    -- Wait for SIGTERM
    takeMVar stopFlag
    putStrLn "Received SIGTERM, shutting down gracefully..."
    killThread acceptThread
    close sock
    removeFile socketPath
    -- Stop proxy on shutdown
    eEnv <- Env.loadEnv environment
    case eEnv of
        Right envVal -> stop (Env.ip envVal) (Env.rawMode envVal)
        Left _ -> pure ()


-- | Accept loop
acceptLoop :: Socket -> MVar () -> FilePath -> String -> IO ()
acceptLoop sock stopFlag config env =
  race serverWait stopWait >> pure ()
  where
    serverWait = forever $ do
      (conn, _) <- accept sock
      forkIO $ handleClient conn config env

    stopWait = readMVar stopFlag   -- Blocks until someone puts something

-- | Handle a single client connection
handleClient :: Socket -> FilePath -> String -> IO ()
handleClient sock configFile environment = do
    h <- socketToHandle sock ReadWriteMode
    hSetBuffering h LineBuffering
    cmdLine <- hGetLine h
    let cmdWords = words cmdLine

    case cmdWords of
        ("start":nameParts) -> do
            let mName = if null nameParts then Nothing else Just (unwords nameParts)

            -- Parse YAML config
            yamlCfg <- parseConfigFile configFile
            case yamlCfg of
                Right cfg -> case mName of
                    -- Select entry by name
                    Just cfgNameStr ->
                      case find (\e -> T.strip (name e) == T.pack cfgNameStr) cfg of
                          Just entry  -> do
                            hPutStrLn h $ "Found " ++ cfgNameStr
                            e <- Env.loadEnv environment
                            case e of
                                Right envVal -> do
                                    stop (Env.ip envVal) (Env.rawMode envVal)
                                    hPutStrLn h "Clearing previous rules"
                                Left _ -> hPutStrLn h "No previous rules detected"
                            ip <- startEntry entry environment
                            hPutStrLn h ("started on " ++ ip)
                          Nothing -> do
                              hPutStrLn h $ "No config entry named: " ++ cfgNameStr
                    Nothing ->
                        case cfg of
                            (entry:_) -> do
                                hPutStrLn h "Starting default server"
                                ip <- startEntry entry environment
                                hPutStrLn h ("Started on " ++ ip)
                            [] ->
                                hPutStrLn h "YAML config empty"

                Left err -> do
                    hPutStrLn h ("No YAML config: " ++ err)

        ["stop"] -> do
            e <- Env.loadEnv environment
            case e of
                Right envVal -> do
                    stop (Env.ip envVal) (Env.rawMode envVal)
                    hPutStrLn h "Stopped"
                Left _ -> hPutStrLn h "No .env loaded, cannot stop"

        ["status"] -> do
            e <- Env.loadEnv environment
            case e of
                 Right envVal -> do
                     hPutStrLn h ("Currently using the " ++ (Env.ip envVal) ++ " server")

        _ -> hPutStrLn h "unknown command"

    hClose h

startEntry :: ConfigEntry -> String -> IO String
startEntry entry environment = do
    ip <- waitForIP (T.unpack $ nameserver entry) (T.unpack $ domain entry)
    let raw = raw_mode entry >>= \rm -> textToRawMode rm (port entry)
    Env.writeEnv Env{Env.ip = ip, rawMode = raw} (T.unpack <$> udp2raw_password entry) (T.unpack <$> udpspeeder_password entry) environment
    start ip raw
    pure ip

-- | Start the proxy
start :: String -> Maybe RawMode -> IO ()
start ip raw = do
    applyRules
    case raw of
        Nothing          -> pure()
        Just ICMP        -> applyICMP ip
        Just (UDP p )    -> applyUDP ip p
        Just (FakeTCP p) -> applyFakeTCP ip p

-- | Stop the proxy
stop :: String -> Maybe RawMode -> IO ()
stop ip raw = do
    clearRules
    case raw of
        Nothing          -> pure()
        Just ICMP        -> clearICMP ip
        Just (UDP p)     -> clearUDP ip p
        Just (FakeTCP p) -> clearFakeTCP ip p