{- | Compile a checked FWL program to nftables JSON using Aeson.
     All policies (Filter and NAT) go into one table named by Config.
     Layer stripping: Frame patterns that omit Ether compile identically
     to those that include it — the compiler inserts protocol matches
     from whatever constructor the user wrote.
-}
module FWL.Compile
  ( compileProgram
  , compileToJson
  ) where

import Data.List (intercalate)
import Data.Maybe (mapMaybe)
import qualified Data.Map.Strict as Map
import qualified Data.Aeson as A
import Data.Aeson ((.=), Value(..), object, toJSON)
import qualified Data.Aeson.Key as K
import qualified Data.ByteString.Lazy as BL
import qualified Data.Aeson.Encode.Pretty as Pretty

import FWL.AST

-- ─── Entry points ────────────────────────────────────────────────────────────

-- | Compile an FWL program and render to pretty-printed JSON bytes.
compileToJson :: Program -> BL.ByteString
compileToJson = Pretty.encodePretty . programToValue

-- | Compile an FWL program to an Aeson Value (the nftables JSON schema).
programToValue :: Program -> Value
programToValue prog@(Program cfg decls) =
  object [ "nftables" .= toJSON (metainfo : tableObj : chainObjs ++ mapObjs ++ ruleObjs) ]
  where
    env       = buildEnv decls
    tbl       = configTable cfg

    metainfo  = object [ "metainfo" .= object [ "json_schema_version" .= (1 :: Int) ] ]
    tableObj  = object [ "table"    .= tableValue tbl ]

    policies  = [ (n, pm, ab) | DPolicy n _ pm ab <- decls ]
    chainObjs = map (\(n, pm, _)  -> chainDeclValue tbl n pm) policies
    ruleObjs  = concatMap (\(n, _, ab) -> concatMap (armToRuleValues env tbl n) ab) policies

    letDecls  = [ (n, t, e) | DLet n t e <- decls ]
    mapObjs   = mapMaybe (\(n, _, e) -> letToMapValue tbl n e) letDecls

-- ─── Table / Chain declarations ──────────────────────────────────────────────

tableValue :: String -> Value
tableValue tbl = object
  [ "family" .= ("inet" :: String)
  , "name"   .= tbl
  ]

chainDeclValue :: String -> Name -> PolicyMeta -> Value
chainDeclValue tbl n pm = object
  [ "chain" .= object
      [ "family"   .= ("inet" :: String)
      , "table"    .= tbl
      , "name"     .= n
      , "type"     .= chainTypeStr (pmTable pm)
      , "hook"     .= hookStr (pmHook pm)
      , "prio"     .= priorityStr (pmPriority pm)
      , "policy"   .= defaultPolicyStr (pmHook pm)
      ]
  ]

chainTypeStr :: TableName -> String
chainTypeStr TFilter = "filter"
chainTypeStr TNAT    = "nat"

hookStr :: Hook -> String
hookStr HInput       = "input"
hookStr HForward     = "forward"
hookStr HOutput      = "output"
hookStr HPrerouting  = "prerouting"
hookStr HPostrouting = "postrouting"

priorityStr :: Priority -> String
priorityStr PFilter  = "filter"
priorityStr PDstNat  = "dstnat"
priorityStr PSrcNat  = "srcnat"
priorityStr PMangle  = "mangle"
priorityStr (PInt n) = show n

-- Input and Forward hooks default to drop; everything else to accept.
defaultPolicyStr :: Hook -> String
defaultPolicyStr HInput   = "drop"
defaultPolicyStr HForward = "drop"
defaultPolicyStr _        = "accept"

-- ─── Arm → Rule objects ──────────────────────────────────────────────────────

-- Each policy arm becomes zero or more nftables rule objects.
-- An arm whose action is Continue compiles to zero rules.
armToRuleValues :: CompileEnv -> String -> Name -> Arm -> [Value]
armToRuleValues env tbl chain (Arm p mg body) =
  case compileAction env body of
    Nothing      -> []   -- Continue: emit nothing
    Just verdict ->
      let patExprs   = compilePat env p
          guardExprs = maybe [] (compileGuard env) mg
          allExprs   = patExprs ++ guardExprs ++ [verdict]
      in  [ object
              [ "rule" .= object
                  [ "family" .= ("inet" :: String)
                  , "table"  .= tbl
                  , "chain"  .= chain
                  , "expr"   .= toJSON allExprs
                  ]
              ]
          ]

-- ─── Pattern → [Value] ───────────────────────────────────────────────────────

type CompileEnv = Map.Map String Decl

buildEnv :: [Decl] -> CompileEnv
buildEnv = foldr (\d m -> Map.insert (declNameOf d) d m) Map.empty
  where
    declNameOf (DInterface n _ _) = n
    declNameOf (DZone n _)        = n
    declNameOf (DPattern n _ _)   = n
    declNameOf (DFlow n _)        = n
    declNameOf (DRule n _ _)      = n
    declNameOf (DPolicy n _ _ _)  = n
    declNameOf (DLet n _ _)       = n
    declNameOf (DImport n _ _)    = n

compilePat :: CompileEnv -> Pat -> [Value]
compilePat _   PWild             = []
compilePat _   (PVar _)          = []
compilePat env (PNamed n)        = expandNamedPat env n
compilePat env (PFrame mp inner) =
  maybe [] (compilePathPat env) mp ++ compilePat env inner
compilePat env (PCtor n ps)      = compileCtorPat env n ps
compilePat _   (PRecord n fs)    = compileRecordPat n fs
compilePat env (PTuple ps)       = concatMap (compilePat env) ps
compilePat _   (PBytes _)        = []   -- handled by flow/ct mark (future)

-- Named patterns are inlined at compile time.
expandNamedPat :: CompileEnv -> Name -> [Value]
expandNamedPat env n =
  case Map.lookup n env of
    Just (DPattern _ _ p) -> compilePat env p
    _                     -> []

-- Layer stripping: Ether is transparent; IPv4/IPv6/TCP/UDP/ICMPv6 each emit
-- the appropriate protocol-selector match then recurse into their children.
-- Omitting Ether produces identical output.
compileCtorPat :: CompileEnv -> String -> [Pat] -> [Value]
compileCtorPat env ctor ps = case ctor of
  "Ether"  -> children                           -- transparent layer
  "IPv4"   -> matchMeta "nfproto" "ipv4"   : children
  "IPv6"   -> matchMeta "nfproto" "ipv6"   : children
  "TCP"    -> matchPayload "th" "protocol" "tcp"   : children
  "UDP"    -> matchPayload "th" "protocol" "udp"   : children
  "ICMPv6" -> matchPayload "ip6" "nexthdr" "ipv6-icmp" : children
  "ICMP"   -> matchPayload "ip"  "protocol" "icmp" : children
  _        -> children
  where
    children = concatMap (compilePat env) ps

-- Record patterns emit field equality matches, e.g. tcp { dport = :22 }.
compileRecordPat :: String -> [FieldPat] -> [Value]
compileRecordPat proto = mapMaybe go
  where
    go (FPEq field lit) = Just (matchPayload proto field (renderLit lit))
    go _                = Nothing

-- Path patterns (iif/oif).
compilePathPat :: CompileEnv -> PathPat -> [Value]
compilePathPat _ (PathPat ms md) =
  maybe [] (compileEndpoint "iifname") ms ++
  maybe [] (compileEndpoint "oifname") md

compileEndpoint :: String -> EndpointPat -> [Value]
compileEndpoint _   EPWild          = []
compileEndpoint dir (EPName n)      = [matchMeta dir n]
compileEndpoint dir (EPMember _ z)  = [matchInSet (metaVal dir) [z]]
  -- zone membership: for MVP we emit the zone name as a set element.
  -- A later pass would expand zones to their member interface names.

-- ─── Guard → [Value] ─────────────────────────────────────────────────────────

compileGuard :: CompileEnv -> Expr -> [Value]
compileGuard env (EInfix OpAnd l r)   = compileGuard env l ++ compileGuard env r
compileGuard _   (EInfix OpIn  l r)   = [compileInExpr l r]
compileGuard _   (EInfix OpEq  l r)   = [matchExpr "==" (exprVal l) (exprVal r)]
compileGuard _   (EInfix OpNeq l r)   = [matchExpr "!=" (exprVal l) (exprVal r)]
compileGuard _   _                    = []

compileInExpr :: Expr -> Expr -> Value
-- ct.state in { Established, Related }
compileInExpr (EQual ["ct","state"])  (ESet vs) = ctMatch "state"  vs
compileInExpr (EQual ["ct","status"]) (ESet vs) = ctMatch "status" vs
-- generic set membership
compileInExpr l (ESet vs) = matchExpr "in" (exprVal l) (setVal (map exprToStr vs))
compileInExpr l r         = matchExpr "==" (exprVal l) (exprVal r)

ctMatch :: String -> [Expr] -> Value
ctMatch key vs = matchExpr "in"
  (object ["ct" .= object ["key" .= key]])
  (setVal (map exprToStr vs))

-- ─── Action → Maybe Value (Nothing = Continue = no rule) ─────────────────────

compileAction :: CompileEnv -> Expr -> Maybe Value
compileAction _   (EVar "Allow")      = Just (object ["accept"     .= Null])
compileAction _   (EVar "Drop")       = Just (object ["drop"       .= Null])
compileAction _   (EVar "Continue")   = Nothing
compileAction _   (EVar "Masquerade") = Just (object ["masquerade" .= Null])
compileAction _   (EApp (EVar "DNAT") arg) =
  Just $ object ["dnat" .= object ["addr" .= exprToStr arg]]
compileAction _   (EApp (EVar "DNATMap") arg) =
  Just $ object ["dnat" .= object ["addr" .= object
    ["map" .= object ["key" .= object ["concat" .= Array mempty]
                     ,"data" .= exprToStr arg]]]]
-- Rule invocation → jump
compileAction env (EApp (EVar rn) _) =
  case Map.lookup rn env of
    Just (DRule _ _ _) -> Just $ object ["jump" .= object ["target" .= rn]]
    _                  -> Just (object ["accept" .= Null])
compileAction _ _ = Just (object ["accept" .= Null])

-- ─── Let → Map object ────────────────────────────────────────────────────────

letToMapValue :: String -> Name -> Expr -> Maybe Value
letToMapValue tbl n (EMap entries) = Just $ object
  [ "map" .= object
      [ "family"   .= ("inet" :: String)
      , "table"    .= tbl
      , "name"     .= n
      , "type"     .= ("inetproto . inetservice" :: String)
      , "map"      .= ("ipv4_addr . inetservice" :: String)
      , "elem"     .= toJSON (map renderMapElem entries)
      ]
  ]
letToMapValue _ _ _ = Nothing

renderMapElem :: (Expr, Expr) -> Value
renderMapElem (k, v) = toJSON
  [ object ["concat" .= toJSON [exprToStr k]]
  , exprToStr v
  ]

-- ─── Aeson building blocks ───────────────────────────────────────────────────

-- { "match": { "op": op, "left": left, "right": right } }
matchExpr :: String -> Value -> Value -> Value
matchExpr op l r = object
  [ "match" .= object
      [ "op"    .= op
      , "left"  .= l
      , "right" .= r
      ]
  ]

matchMeta :: String -> String -> Value
matchMeta key val = matchExpr "==" (metaVal key) (A.String (strText val))

matchPayload :: String -> String -> String -> Value
matchPayload proto field val =
  matchExpr "==" (payloadVal proto field) (A.String (strText val))

matchInSet :: Value -> [String] -> Value
matchInSet lhs vals =
  matchExpr "in" lhs (setVal vals)

metaVal :: String -> Value
metaVal key = object ["meta" .= object ["key" .= key]]

payloadVal :: String -> String -> Value
payloadVal proto field =
  object ["payload" .= object ["protocol" .= proto, "field" .= field]]

setVal :: [String] -> Value
setVal vs = object ["set" .= toJSON vs]

-- ─── Expression → Value helpers ──────────────────────────────────────────────

exprVal :: Expr -> Value
exprVal (EQual [p, f])          = payloadVal p f
exprVal (EQual ["ct", k])       = object ["ct" .= object ["key" .= k]]
exprVal (EVar n)                = metaVal n
exprVal (ELit l)                = A.String (strText (renderLit l))
exprVal (ESet vs)               = setVal (map exprToStr vs)
exprVal e                       = A.String (strText (exprToStr e))

exprToStr :: Expr -> String
exprToStr (EVar n)    = n
exprToStr (ELit l)    = renderLit l
exprToStr (EQual ns)  = intercalate "." ns
exprToStr (ETuple es) = intercalate " . " (map exprToStr es)
exprToStr _           = "_"

strText :: String -> A.Text
strText = \s -> read (show s)   -- simple String→Text without extra dep

renderLit :: Literal -> String
renderLit (LInt n)             = show n
renderLit (LString s)          = s
renderLit (LBool True)         = "true"
renderLit (LBool False)        = "false"
renderLit (LIPv4 (a,b,c,d))    =
  show a++"."++show b++"."++show c++"."++show d
renderLit (LIPv6 _)             = "::1"
renderLit (LCIDR ip p)          = renderLit ip ++ "/" ++ show p
renderLit (LPort p)             = show p
renderLit (LDuration n Seconds) = show n
renderLit (LDuration n _)       = show n
renderLit (LHex b)              = show b

-- Data.Aeson.Key helper (aeson >= 2.0 uses Key, not Text, for object keys)
(.=) :: A.ToJSON v => String -> v -> A.Pair
k .= v = (K.fromString k, toJSON v)