{-# LANGUAGE OverloadedStrings #-}

module Env
  ( Env(..)
  , loadEnv
  , writeEnv
  ) where

import System.IO
import Control.Exception (catch, IOException)
import Data.List (isPrefixOf)
import Data.Maybe (fromMaybe)
import Data.Char (toUpper)
import Text.Read (readMaybe)

import Udp2Raw  -- Import the RawMode type (ICMP, UDP, FakeTCP)
import System.IO.Temp (withTempFile)
import System.Directory (renameFile)
import System.FilePath (takeDirectory)

-------------------------------------------------------------------------------
-- | Environment configuration
-------------------------------------------------------------------------------
data Env = Env
  { ip      :: String
  , rawMode :: Maybe RawMode
  } deriving (Show, Eq)

write ::  String -> --filename
          String -> --content
          IO ()
write path content = do
  let dir = takeDirectory path
  withTempFile dir "env.tmp" $ \tmpPath tmpHandle -> do
    hPutStrLn tmpHandle (content)
    hClose tmpHandle
    renameFile tmpPath path  -- atomic replacement


writeEnv :: Env ->
            Maybe String -> --UDP2RAW_PWD
            Maybe String -> --UDP2SPEEDER_PWD
            String -> --filename
            IO ()
writeEnv env udp2raw_pwd udpspeeder_pwd file = do
    write file (
        "SERVER_IP=" ++ ip env ++ "\n"
        ++ "UDP2RAW_PWD=" ++ fromMaybe "" udp2raw_pwd ++ "\n"
        ++ "UDPSPEEDER_PWD=" ++ fromMaybe "" udpspeeder_pwd ++ "\n"
        ++ "RAW_MODE=" ++ maybe "" rawModeToArg (rawMode env) ++ "\n"
        ++ "PORT=" ++ show (fromMaybe 443 (rawMode env >>= getPort)))


-------------------------------------------------------------------------------
-- | Load and parse the .env file
-------------------------------------------------------------------------------
loadEnv :: FilePath -> IO (Either String Env)
loadEnv path = catch (do
    content <- readFile path
    let kvs = parseEnv content
        mIP = lookup "SERVER_IP" kvs
        mMode = lookup "RAW_MODE" kvs
        mPort = lookup "PORT" kvs
    case (mIP, mMode) of
      (Just ipStr, Just modeStr)->pure $ Right $ Env ipStr (parseRawMode modeStr mPort)
      (Just ipStr, _)-> pure $ Right $ Env ipStr Nothing
      _ -> pure $ Left "Missing SERVER_IP or RAW_MODE in .env"
  ) handleReadError
  where
    handleReadError :: IOException -> IO (Either String Env)
    handleReadError e = pure $ Left $ "Error reading .env: " ++ show e

-------------------------------------------------------------------------------
-- | Parse RAW_MODE + optional PORT into RawMode
-------------------------------------------------------------------------------
parseRawMode :: String -> Maybe String -> Maybe RawMode
parseRawMode modeStr mPort =
  let mode = map toUpper modeStr
      port = mPort >>= readMaybe
  in case mode of
       "ICMP"    -> Just ICMP
       "UDP"     -> UDP <$> port         -- mandatory port
       "FAKETCP" -> FakeTCP <$> port     -- mandatory port
       _         -> Nothing

-------------------------------------------------------------------------------
-- | Parse .env content into key-value pairs
-------------------------------------------------------------------------------
parseEnv :: String -> [(String, String)]
parseEnv = mapMaybe parseLine . lines
  where
    parseLine line =
      case break (== '=') (trim line) of
        (key, '=':value)
          | not (null key) && not (isComment key) ->
              Just (trim key, trim value)
        _ -> Nothing

    isComment ('#':_) = True
    isComment _       = False

    trim = f . f
      where f = reverse . dropWhile (`elem` [' ', '\t', '\r'])

-- | Helper: mapMaybe for base < 4.8
mapMaybe :: (a -> Maybe b) -> [a] -> [b]
mapMaybe f = foldr (\x acc -> case f x of
                               Just y  -> y:acc
                               Nothing -> acc) []