module FWL.Parser
  ( parseProgram
  , parseFile
  ) where

import Control.Monad (void, when)
import Data.Bits ((.&.), (.|.), shiftL)
import Data.List (foldl')
import Data.Word (Word8)
import Numeric (readHex)
import Text.Parsec
import Text.Parsec.String (Parser)
import Data.Functor.Identity (Identity)
import qualified Text.Parsec.Expr as Ex

import FWL.AST
import FWL.Lexer
import Data.Char (isUpper)

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

parseProgram :: String -> String -> Either ParseError Program
parseProgram src input = parse program src input

parseFile :: FilePath -> IO (Either ParseError Program)
parseFile fp = parseProgram fp <$> readFile fp

-- ─── Top-level ───────────────────────────────────────────────────────────────

program :: Parser Program
program = do
  whiteSpace
  cfg  <- option defaultConfig configBlock
  ds   <- many decl
  eof
  return (Program cfg ds)

configBlock :: Parser Config
configBlock = do
  reserved "config"
  props <- braces (endBy configProp semi)
  optional semi
  return $ foldr applyProp defaultConfig props
  where
    applyProp ("table", v) c = c { configTable = v }
    applyProp _            c = c

configProp :: Parser (String, String)
configProp = do
  n <- identifier  -- "table" is no longer reserved
  reservedOp "="
  v <- stringLit
  return (n, v)

-- ─── Declarations ────────────────────────────────────────────────────────────

decl :: Parser Decl
decl =  interfaceDecl
    <|> zoneDecl
    <|> importDecl
    <|> letDecl
    <|> patternDecl
    <|> flowDecl
    <|> ruleDecl
    <|> policyDecl
    <|> portforwardDecl
    <|> masqueradeDecl

interfaceDecl :: Parser Decl
interfaceDecl = do
  reserved "interface"
  n  <- identifier
  reservedOp ":"
  k  <- ifaceKind
  ps <- braces (endBy ifaceProp semi)
  _ <- semi
  return (DInterface n k ps)

ifaceKind :: Parser IfaceKind
ifaceKind =  (reserved "WAN"       >> return IWan)
         <|> (reserved "LAN"       >> return ILan)
         <|> (reserved "WireGuard" >> return IWireGuard)
         <|> (IUser <$> identifier)

ifaceProp :: Parser IfaceProp
ifaceProp =  (reserved "dynamic" >> return IPDynamic)
         <|> (reserved "cidr4" >> reservedOp "=" >> IPCidr4 <$> cidrSet)
         <|> (reserved "cidr6" >> reservedOp "=" >> IPCidr6 <$> cidrSet)

cidrSet :: Parser [CIDR]
cidrSet = braces (commaSep1 cidrLit)

zoneDecl :: Parser Decl
zoneDecl = do
  reserved "zone"
  n  <- identifier
  reservedOp "="
  ns <- braces (commaSep1 identifier)
  _ <- semi
  return (DZone n ns)

importDecl :: Parser Decl
importDecl = do
  reserved "import"
  n <- identifier
  reservedOp ":"
  t <- typeP
  reserved "from"
  s <- stringLit
  _ <- semi
  return (DImport n t s)

letDecl :: Parser Decl
letDecl = do
  reserved "let"
  n <- identifier
  reservedOp ":"
  t <- typeP
  reservedOp "="
  e <- expr
  _ <- semi
  return (DLet n t e)

patternDecl :: Parser Decl
patternDecl = do
  reserved "pattern"
  n <- identifier
  reservedOp ":"
  t <- typeP
  reservedOp "="
  p <- pat
  _ <- semi
  return (DPattern n t p)

flowDecl :: Parser Decl
flowDecl = do
  reserved "flow"
  n <- identifier
  reservedOp ":"
  reserved "FlowPattern"
  reservedOp "="
  f <- flowExpr
  _ <- semi
  return (DFlow n f)

ruleDecl :: Parser Decl
ruleDecl = do
  reserved "rule"
  n <- identifier
  reservedOp ":"
  t <- typeP
  reservedOp "="
  e <- expr
  _ <- semi
  return (DRule n t e)

policyDecl :: Parser Decl
policyDecl = do
  reserved "policy"
  n  <- identifier
  reservedOp ":"
  t  <- typeP
  reserved "hook"
  h  <- hookP
  mp <- optionMaybe (reserved "priority" >> priorityP)
  let tb = hookDefaultTable h
      pr = maybe (hookDefaultPriority h) id mp
  reservedOp "="
  ab <- armBlock
  _  <- semi
  return (DPolicy n t (PolicyMeta h tb pr) ab)

-- | Infer table from hook
hookDefaultTable :: Hook -> TableName
hookDefaultTable HInput       = TFilter
hookDefaultTable HForward     = TFilter
hookDefaultTable HOutput      = TFilter
hookDefaultTable HPrerouting  = TNAT
hookDefaultTable HPostrouting = TNAT

-- | Default priority per hook
hookDefaultPriority :: Hook -> Priority
hookDefaultPriority HInput       = pFilter
hookDefaultPriority HForward     = pFilter
hookDefaultPriority HOutput      = pFilter
hookDefaultPriority HPrerouting  = pDstNat
hookDefaultPriority HPostrouting = pSrcNat

hookP :: Parser Hook
hookP =  (reserved "Input"       >> return HInput)
     <|> (reserved "Forward"     >> return HForward)
     <|> (reserved "Output"      >> return HOutput)
     <|> (reserved "Prerouting"  >> return HPrerouting)
     <|> (reserved "Postrouting" >> return HPostrouting)

-- portforward <name> on <iface> via <MapType> = { entries };
portforwardDecl :: Parser Decl
portforwardDecl = do
  reserved "portforward"
  n    <- identifier
  reserved "on"
  iface <- identifier
  reserved "via"
  t    <- typeP
  reservedOp "="
  entries <- braces (commaSep mapEntry)
  _    <- semi
  return (DPortForward n iface t entries)

-- masquerade <name> on <iface> src <set-name>;
masqueradeDecl :: Parser Decl
masqueradeDecl = do
  reserved "masquerade"
  n     <- identifier
  reserved "on"
  iface <- identifier
  reserved "src"
  srcSet <- identifier
  _     <- semi
  return (DMasquerade n iface srcSet)

priorityP :: Parser Priority
priorityP
  =  (reserved "Filter"  >> return pFilter)
 <|> (reserved "DstNat"  >> return pDstNat)
 <|> (reserved "SrcNat"  >> return pSrcNat)
 <|> (reserved "Mangle"  >> return pMangle)
 <|> (reserved "Raw"     >> return pRaw)
 <|> (reserved "ConnTrack" >> return pConnTrack)
 <|> (Priority . fromIntegral <$> integerP)
  where
    -- Accept optional leading minus for negative priorities
    integerP = do
      neg <- option 1 (char '-' >> return (-1))
      n   <- natural
      whiteSpace
      return (neg * fromIntegral n)

-- ─── Arm blocks ──────────────────────────────────────────────────────────────

armBlock :: Parser ArmBlock
armBlock = braces (many arm)

arm :: Parser Arm
arm = do
  _ <- symbol "|"
  p  <- pat
  g  <- optionMaybe (reserved "if" >> expr)
  reservedOp "->"
  e  <- expr
  _ <- semi
  return (Arm p g e)

-- ─── Patterns ────────────────────────────────────────────────────────────────

pat :: Parser Pat
pat = Ex.buildExpressionParser patTable patAtom <?> "pattern"
  where
    patTable = [ [Ex.Infix (reservedOp "|" >> return POr) Ex.AssocLeft] ]

patAtom :: Parser Pat
patAtom =  wildcardPat
   <|> try framePat
   <|> try tuplePat
   <|> bytesPat
   <|> try recordPat
   <|> try namedOrCtorPat

wildcardPat :: Parser Pat
wildcardPat = symbol "_" >> return PWild

-- Frame(...) — optional path then inner pattern
-- Layer stripping: if the inner pattern is not Ether/IPv4/IPv6/etc the
-- type-checker will peel outer layers automatically. Parser just stores
-- whatever the user wrote.
framePat :: Parser Pat
framePat = do
  reserved "Frame"
  (mp, inner) <- parens frameArgs
  return (PFrame mp inner)

frameArgs :: Parser (Maybe PathPat, Pat)
frameArgs = try withPath <|> withoutPath
  where
    withPath = do
      pp <- pathPat
      _ <- comma
      inner <- pat
      return (Just pp, inner)
    withoutPath = do
      inner <- pat
      return (Nothing, inner)

pathPat :: Parser PathPat
pathPat = do
  src <- optionMaybe (try endpointPat)
  dst <- optionMaybe (try (reservedOp "->" >> endpointPat))
  case (src, dst) of
    (Nothing, Nothing) -> fail "empty path pattern"
    _                  -> return (PathPat src dst)

endpointPat :: Parser EndpointPat
endpointPat
  =  (symbol "_" >> return EPWild)
 <|> try (do n <- identifier
             memberOp
             z <- identifier
             return (EPMember n z))
 <|> (EPName <$> identifier)

memberOp :: Parser ()
memberOp = (reservedOp "∈" <|> reserved "in") >> return ()

tuplePat :: Parser Pat
tuplePat = do
  ps <- parens (commaSep2 pat)
  return (PTuple ps)

commaSep2 :: Parser a -> Parser [a]
commaSep2 p = do
  x  <- p
  _ <- comma
  xs <- commaSep1 p
  return (x:xs)

bytesPat :: Parser Pat
bytesPat = brackets (PBytes <$> many byteElem)

byteElem :: Parser ByteElem
byteElem
  =  try (symbol "_*" >> return BEWildStar)
 <|> try (symbol "_"  >> return BEWild)
 <|> (BEHex <$> hexByte)

hexByte :: Parser Word8
hexByte = do
  void (string "0x")
  h1 <- hexDigit
  h2 <- hexDigit
  whiteSpace
  case (readHex [h1,h2] :: [(Integer, String)]) of
    [(v,"")] -> return (fromIntegral v)
    _        -> fail "invalid hex byte"

-- Record pattern: ident { fields }
recordPat :: Parser Pat
recordPat = do
  n  <- identifier
  fs <- braces (commaSep fieldPat)
  return (PRecord n fs)

fieldPat :: Parser FieldPat
fieldPat = do
  n <- identifier
  try (reservedOp "=" >> FPEq n <$> fieldLiteral)
    <|> try (reserved "as" >> FPAs n <$> identifier)
    <|> return (FPBind n)

-- Port literals (:22) are valid in record field position as well as plain literals.
fieldLiteral :: Parser Literal
fieldLiteral = try portLit <|> literal
  where
    portLit = do
      void (char ':')
      n <- fromIntegral <$> natural
      return (LPort n)

namedOrCtorPat :: Parser Pat
namedOrCtorPat = do
  n  <- identifier
  args <- optionMaybe (try (parens (commaSep pat)))
  case args of
    Nothing -> return $ if null n then PWild else if isUpper (head n) then PNamed n else PVar n
    Just ps -> return (PCtor  n ps)

-- ─── Flow expressions ────────────────────────────────────────────────────────

flowExpr :: Parser FlowExpr
flowExpr = do
  first <- FAtom <$> identifier
  rest  <- many (reservedOp "." >> identifier)
  mw    <- optionMaybe (reserved "within" >> durationLit)
  let chain = buildSeq (first : map FAtom rest)
  return $ case mw of
    Nothing -> chain
    Just w  -> attach w chain
  where
    buildSeq [x]    = x
    buildSeq (x:xs) = FSeq x (buildSeq xs) Nothing
    buildSeq []     = error "impossible"

    attach w (FSeq a b _) = FSeq a b (Just w)
    attach w x            = FSeq x x (Just w)

durationLit :: Parser Duration
durationLit = do
  n <- fromIntegral <$> natural
  u <-  (char 's' >> return Seconds)
    <|> (string "ms" >> return Millis)
    <|> (char 'm' >> return Minutes)
    <|> (char 'h' >> return Hours)
  whiteSpace
  return (n, u)

-- ─── Types ───────────────────────────────────────────────────────────────────

typeP :: Parser Type
typeP = do
  t <- baseType
  option t (reservedOp "->" >> TFun t <$> typeP)

baseType :: Parser Type
baseType
  =  effectType
 <|> try tupleTy
 <|> simpleTy

effectType :: Parser Type
effectType = do
  effs <- angles (commaSep identifier)
  t    <- simpleTy
  return (TEffect effs t)

tupleTy :: Parser Type
tupleTy = TTuple <$> parens (commaSep2 typeP)

simpleTy :: Parser Type
simpleTy = do
  n    <- identifier
  args <- option [] (angles (commaSep typeP))
  return (TName n args)

-- ─── Expressions ─────────────────────────────────────────────────────────────

expr :: Parser Expr
expr =  lamExpr
    <|> ifExpr
    <|> doExpr
    <|> caseExpr
    <|> letExpr
    <|> infixExpr

lamExpr :: Parser Expr
lamExpr = do
  reservedOp "\\"
  n <- identifier
  reservedOp "->"
  e <- expr
  return (ELam n e)

ifExpr :: Parser Expr
ifExpr = do
  reserved "if"
  c <- expr
  reserved "then"
  t <- expr
  reserved "else"
  f <- expr
  return (EIf c t f)

doExpr :: Parser Expr
doExpr = reserved "do" >> braces (EDo <$> semiSep doStmt)

doStmt :: Parser DoStmt
doStmt = try bindStmt <|> (DSExpr <$> expr)

bindStmt :: Parser DoStmt
bindStmt = do
  n <- identifier
  reservedOp "<-"
  e <- expr
  return (DSBind n e)

caseExpr :: Parser Expr
caseExpr = do
  reserved "case"
  e  <- expr
  reserved "of"
  ab <- armBlock
  return (ECase e ab)

letExpr :: Parser Expr
letExpr = do
  reserved "let"
  n <- identifier
  reservedOp "="
  e1 <- expr
  reserved "in"
  e2 <- expr
  return (ELet n e1 e2)

-- Operator table for infix expressions
infixExpr :: Parser Expr
infixExpr = Ex.buildExpressionParser opTable appExpr

opTable :: Ex.OperatorTable String () Identity Expr
opTable =
  [ [ prefix "!" ENot ]
  , [ infixL "==" OpEq, infixL "!=" OpNeq
    , infixL "<"  OpLt, infixL "<=" OpLte
    , infixL ">"  OpGt, infixL ">=" OpGte
    , infixIn ]
  , [ infixR "&&" OpAnd ]
  , [ infixR "||" OpOr  ]
  , [ infixR "++" OpConcat ]
  , [ infixL ">>=" OpBind ]
  , [ infixL ">>"  OpThen ]
  ]
  where
    prefix op f = Ex.Prefix  (reservedOp op >> return f)
    infixL op c = Ex.Infix   (reservedOp op >> return (EInfix c)) Ex.AssocLeft
    infixR op c = Ex.Infix   (reservedOp op >> return (EInfix c)) Ex.AssocRight
    infixIn     = Ex.Infix
      ((memberOp <|> reserved "in") >> return (EInfix OpIn))
      Ex.AssocNone

appExpr :: Parser Expr
appExpr = do
  f    <- atom
  args <- many atom
  return (foldl EApp f args)

atom :: Parser Expr
atom
  =  try performExpr
 <|> try mapLit
 <|> try setLit
 <|> try tupleLit
 <|> try (parens expr)
 <|> try litExpr
 <|> try portExpr
 <|> qualNameExpr

performExpr :: Parser Expr
performExpr = do
  reserved "perform"
  parts <- sepBy1 identifier dot
  args  <- parens (commaSep expr)
  return (EPerform parts args)

qualNameExpr :: Parser Expr
qualNameExpr = do
  parts <- sepBy1 identifier (try (dot <* notFollowedBy digit))
  case parts of
    [n] -> return (EVar n)
    ns  -> return (EQual ns)

litExpr :: Parser Expr
litExpr = ELit <$> literal

portExpr :: Parser Expr
portExpr = do
  void (char ':')
  n <- fromIntegral <$> natural
  return (ELit (LPort n))

tupleLit :: Parser Expr
tupleLit = ETuple <$> parens (commaSep2 expr)

setLit :: Parser Expr
setLit = braces $ do
  items <- commaSep expr
  return (ESet items)

-- map literal: { expr -> expr, ... }
mapLit :: Parser Expr
mapLit = braces $ do
  entries <- commaSep1 mapEntry
  return (EMap entries)

mapEntry :: Parser (Expr, Expr)
mapEntry = do
  k <- expr
  reservedOp "->"
  v <- expr
  return (k, v)

-- ─── Literals ────────────────────────────────────────────────────────────────

literal :: Parser Literal
literal
  =  try ipOrCidrLit
 <|> try hexLit
 <|> try (LBool True  <$ reserved "true")
 <|> try (LBool False <$ reserved "false")
 <|> try (LString <$> stringLit)
 <|> try (LInt . fromIntegral <$> natural)

hexLit :: Parser Literal
hexLit = LHex <$> hexByte

-- ─── IP / CIDR parsing ───────────────────────────────────────────────────────

-- | Parse an IPv4 or IPv6 address, optionally followed by /prefix.
-- Tries IPv6 first (it can start with hex digits too), then IPv4.
ipOrCidrLit :: Parser Literal
ipOrCidrLit = do
  ip <- try ipv6Lit <|> ipv4Lit_
  mPrefix <- optionMaybe (char '/' >> fromIntegral <$> natural)
  whiteSpace
  return $ case mPrefix of
    Nothing -> ip
    Just p  -> LCIDR ip p

-- | IPv4: four decimal octets separated by dots → LIP IPv4 (32-bit Integer)
ipv4Lit_ :: Parser Literal
ipv4Lit_ = do
  a <- octet
  void (char '.')
  b <- octet
  void (char '.')
  c <- octet
  void (char '.')
  d <- octet
  return $ LIP IPv4
    (  fromIntegral a `shiftL` 24
   .|. fromIntegral b `shiftL` 16
   .|. fromIntegral c `shiftL` 8
   .|. fromIntegral d)
  where
    octet = do
      n <- fromIntegral <$> natural
      if n > 255 then fail "octet out of range" else return n

-- | IPv6: full notation, :: abbreviation, and optional embedded IPv4.
-- Stores as LIP IPv6 (128-bit Integer).
ipv6Lit :: Parser Literal
ipv6Lit = do
  (left, hasDbl, right) <- ipv6Groups
  let missing   = 8 - length left - length right
  when (missing < 0) $ fail "too many groups in IPv6 address"
  when (not hasDbl && missing /= 0) $ fail "invalid IPv6 address (must have 8 groups or use ::)"
  let groups    = left ++ replicate missing 0 ++ right
  when (length groups /= 8) $ fail "invalid IPv6 address"
  let val = foldl' (\acc g -> (acc `shiftL` 16) .|. fromIntegral g) (0::Integer) groups
  return (LIP IPv6 val)

-- Returns (left-of-::, has_dbl_colon, right-of-::).
-- If no :: present, left has all 8 groups and right is empty.
ipv6Groups :: Parser ([Int], Bool, [Int])
ipv6Groups = do
  -- must start with a hex digit or ':' (for ::)
  ahead <- lookAhead (hexDigit <|> char ':')
  case ahead of
    ':' -> do
      void (string "::")
      right <- ipv6RightGroups
      return ([], True, right)
    _   -> do
      left <- ipv6LeftGroups
      mDbl <- optionMaybe (try (string "::"))
      case mDbl of
        Nothing -> return (left, False, [])
        Just _  -> do
          right <- ipv6RightGroups
          return (left, True, right)

-- Parse a run of hex16:hex16:... stopping before :: or end
ipv6LeftGroups :: Parser [Int]
ipv6LeftGroups = do
  first <- hex16
  rest  <- many (try (char ':' >> notFollowedBy (char ':') >> hex16))
  return (first : rest)

-- Parse groups to the right of ::, including optional embedded IPv4
ipv6RightGroups :: Parser [Int]
ipv6RightGroups = option [] $
  try ipv4EmbeddedGroups <|> ipv6LeftGroups

-- IPv4-mapped groups: e.g. ffff:192.168.1.1 -> [0xffff, 0xc0a8, 0x0101]
ipv4EmbeddedGroups :: Parser [Int]
ipv4EmbeddedGroups = do
  prefix <- many (try (hex16 <* char ':' <* lookAhead digit))
  a <- octet_; void (char '.')
  b <- octet_; void (char '.')
  c <- octet_; void (char '.')
  d <- octet_
  let hi = (a `shiftL` 8) .|. b
      lo = (c `shiftL` 8) .|. d
  return (prefix ++ [hi, lo])
  where
    octet_ = do
      n <- fromIntegral <$> natural
      if n > 255 then fail "IPv4 octet out of range" else return n

hex16 :: Parser Int
hex16 = do
  digits <- many1 hexDigit
  case (reads ("0x" ++ digits)) :: [(Int,String)] of
    [(v,"")] -> if v > 0xffff then fail "hex16 out of range" else return v
    _        -> fail "invalid hex group"

cidrLit :: Parser CIDR
cidrLit = do
  l <- ipOrCidrLit
  case l of
    LCIDR ip p -> return (ip, p)
    _          -> fail "expected CIDR notation (address/prefix)"