module ParserTests (tests) where

import Test.Tasty
import Test.Tasty.HUnit

import FWL.AST
import FWL.Util

tests :: TestTree
tests = testGroup "Parser"
  [ interfaceTests
  , zoneTests
  , importTests
  , letTests
  , patternTests
  , flowTests
  , typeTests
  , exprTests
  , policyTests
  , ruleTests
  , configTests
  , errorTests
  ]

-- ─── Interface ───────────────────────────────────────────────────────────────

interfaceTests :: TestTree
interfaceTests = testGroup "interface"
  [ testCase "WAN dynamic" $ do
      p <- parseOk "interface wan : WAN { dynamic; };"
      d <- singleDecl p
      case d of
        DInterface "wan" IWan [IPDynamic] -> return ()
        _ -> assertFailure (show d)

  , testCase "LAN with cidr4" $ do
      p <- parseOk "interface lan : LAN { cidr4 = { 10.0.0.0/8 }; };"
      d <- singleDecl p
      case d of
        DInterface "lan" ILan [IPCidr4 [(ip, 8)]] | ip == ipv4Lit 10 0 0 0 -> return ()
        _ -> assertFailure (show d)

  , testCase "LAN with cidr4 and cidr6" $ do
      p <- parseOk
            "interface lan : LAN { \
            \  cidr4 = { 10.17.1.0/24 }; \
            \  cidr6 = { 192.168.0.0/16 }; \
            \};"
      d <- singleDecl p
      case d of
        DInterface "lan" ILan [IPCidr4 _, IPCidr6 _] -> return ()
        _ -> assertFailure (show d)

  , testCase "WireGuard interface" $ do
      p <- parseOk "interface wg0 : WireGuard {};"
      d <- singleDecl p
      case d of
        DInterface "wg0" IWireGuard [] -> return ()
        _ -> assertFailure (show d)

  , testCase "user-defined kind" $ do
      p <- parseOk "interface eth0 : Bridge {};"
      d <- singleDecl p
      case d of
        DInterface "eth0" (IUser "Bridge") [] -> return ()
        _ -> assertFailure (show d)

  , testCase "multiple CIDRs in set" $ do
      p <- parseOk
            "interface lan : LAN { \
            \  cidr4 = { 10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16 }; \
            \};"
      d <- singleDecl p
      case d of
        DInterface _ _ [IPCidr4 cidrs] -> length cidrs @?= 3
        _ -> assertFailure (show d)
  ]

-- ─── Zone ────────────────────────────────────────────────────────────────────

zoneTests :: TestTree
zoneTests = testGroup "zone"
  [ testCase "single member" $ do
      p <- parseOk "zone trusted = { lan };"
      d <- singleDecl p
      case d of
        DZone "trusted" ["lan"] -> return ()
        _ -> assertFailure (show d)

  , testCase "multiple members" $ do
      p <- parseOk "zone lan_zone = { lan, wg0, vlan10 };"
      d <- singleDecl p
      case d of
        DZone "lan_zone" ["lan","wg0","vlan10"] -> return ()
        _ -> assertFailure (show d)
  ]

-- ─── Import ──────────────────────────────────────────────────────────────────

importTests :: TestTree
importTests = testGroup "import"
  [ testCase "basic import" $ do
      p <- parseOk "import rfc1918 : CIDRSet from \"builtin:rfc1918\";"
      d <- singleDecl p
      case d of
        DImport "rfc1918" (TName "CIDRSet" []) "builtin:rfc1918" -> return ()
        _ -> assertFailure (show d)
  ]

-- ─── Let ─────────────────────────────────────────────────────────────────────

letTests :: TestTree
letTests = testGroup "let"
  [ testCase "simple integer" $ do
      p <- parseOk "let timeout : Int = 30;"
      d <- singleDecl p
      case d of
        DLet "timeout" (TName "Int" []) (ELit (LInt 30)) -> return ()
        _ -> assertFailure (show d)

  , testCase "map literal" $ do
      p <- parseOk
            "let forwards : Map<(Protocol,Port),(IP,Port)> = { \
            \  (tcp, :8080) -> (10.0.0.1, :80) \
            \};"
      d <- singleDecl p
      case d of
        DLet "forwards" _ (EMap [_]) -> return ()
        _ -> assertFailure (show d)

  , testCase "string literal" $ do
      p <- parseOk "let name : String = \"hello\";"
      d <- singleDecl p
      case d of
        DLet "name" _ (ELit (LString "hello")) -> return ()
        _ -> assertFailure (show d)
  ]

-- ─── Pattern ─────────────────────────────────────────────────────────────────

patternTests :: TestTree
patternTests = testGroup "pattern"
  [ testCase "tuple with record field" $ do
      p <- parseOk
            "pattern WGInitiation : (UDPHeader, Bytes) = \
            \  (udp { length = 156 }, [0x01 _*]);"
      d <- singleDecl p
      case d of
        DPattern "WGInitiation" _ (PTuple [PRecord "udp" _, PBytes _]) -> return ()
        _ -> assertFailure (show d)

  , testCase "byte pattern elements" $ do
      p <- parseOk
            "pattern WGResponse : (UDPHeader, Bytes) = \
            \  (udp { length = 100 }, [0x02 _ _*]);"
      d <- singleDecl p
      case d of
        DPattern "WGResponse" _ (PTuple [_, PBytes [BEHex 0x02, BEWild, BEWildStar]]) ->
          return ()
        _ -> assertFailure (show d)

  , testCase "named pattern reference in ctor" $ do
      p <- parseOk
            "pattern Complex : Frame = \
            \  Frame(_, IPv4(ip, WGInitiation));"
      d <- singleDecl p
      case d of
        DPattern "Complex" _ (PFrame (Just _) (PCtor "IPv4" [PVar "ip", PNamed "WGInitiation"])) ->
          return ()
        _ -> assertFailure (show d)

  , testCase "record with field bind" $ do
      p <- parseOk "pattern HasTCP : TCP = tcp { dport };"
      d <- singleDecl p
      case d of
        DPattern "HasTCP" _ (PRecord "tcp" [FPBind "dport"]) -> return ()
        _ -> assertFailure (show d)

  , testCase "record with field equality" $ do
      p <- parseOk "pattern SSH : TCP = tcp { dport = :22 };"
      d <- singleDecl p
      case d of
        DPattern "SSH" _ (PRecord "tcp" [FPEq "dport" (LPort 22)]) -> return ()
        _ -> assertFailure (show d)
  ]

-- ─── Flow ────────────────────────────────────────────────────────────────────

flowTests :: TestTree
flowTests = testGroup "flow"
  [ testCase "two-step sequence with within" $ do
      p <- parseOk
            "flow WireGuardHandshake : FlowPattern = \
            \  WGInitiation . WGResponse within 5s;"
      d <- singleDecl p
      case d of
        DFlow "WireGuardHandshake" (FSeq (FAtom "WGInitiation") (FAtom "WGResponse") (Just (5, Seconds))) ->
          return ()
        _ -> assertFailure (show d)

  , testCase "single atom flow" $ do
      p <- parseOk "flow Simple : FlowPattern = Ping;"
      d <- singleDecl p
      case d of
        DFlow "Simple" (FAtom "Ping") -> return ()
        _ -> assertFailure (show d)

  , testCase "duration in milliseconds" $ do
      p <- parseOk "flow Fast : FlowPattern = A . B within 500ms;"
      d <- singleDecl p
      case d of
        DFlow "Fast" (FSeq _ _ (Just (500, Millis))) -> return ()
        _ -> assertFailure (show d)
  ]

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

typeTests :: TestTree
typeTests = testGroup "types"
  [ testCase "simple name" $ do
      p <- parseOk "let x : Frame = Allow;"
      d <- singleDecl p
      case d of
        DLet _ (TName "Frame" []) _ -> return ()
        _ -> assertFailure (show d)

  , testCase "generic type" $ do
      p <- parseOk "let x : Map<Int, String> = Allow;"
      d <- singleDecl p
      case d of
        DLet _ (TName "Map" [TName "Int" [], TName "String" []]) _ -> return ()
        _ -> assertFailure (show d)

  , testCase "function type" $ do
      p <- parseOk "let x : Frame -> Action = Allow;"
      d <- singleDecl p
      case d of
        DLet _ (TFun (TName "Frame" []) (TName "Action" [])) _ -> return ()
        _ -> assertFailure (show d)

  , testCase "effect type" $ do
      p <- parseOk "let x : <Log, FlowMatch> Action = Allow;"
      d <- singleDecl p
      case d of
        DLet _ (TEffect ["Log","FlowMatch"] (TName "Action" [])) _ -> return ()
        _ -> assertFailure (show d)

  , testCase "tuple type" $ do
      p <- parseOk "let x : (Int, String) = Allow;"
      d <- singleDecl p
      case d of
        DLet _ (TTuple [TName "Int" [], TName "String" []]) _ -> return ()
        _ -> assertFailure (show d)

  , testCase "function with effects" $ do
      p <- parseOk "let x : Frame -> <Log> Action = Allow;"
      d <- singleDecl p
      case d of
        DLet _ (TFun _ (TEffect ["Log"] _)) _ -> return ()
        _ -> assertFailure (show d)
  ]

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

exprTests :: TestTree
exprTests = testGroup "expressions"
  [ testCase "boolean and" $ do
      p <- parseOk "let x : Bool = a && b;"
      d <- singleDecl p
      case d of
        DLet _ _ (EInfix OpAnd (EVar "a") (EVar "b")) -> return ()
        _ -> assertFailure (show d)

  , testCase "set membership with 'in'" $ do
      p <- parseOk "let x : Bool = ct.state in { Established, Related };"
      d <- singleDecl p
      case d of
        DLet _ _ (EInfix OpIn (EQual ["ct","state"]) (ESet _)) -> return ()
        _ -> assertFailure (show d)

  , testCase "equality comparison" $ do
      p <- parseOk "let x : Bool = tcp.dport == :22;"
      d <- singleDecl p
      case d of
        DLet _ _ (EInfix OpEq (EQual ["tcp","dport"]) (ELit (LPort 22))) -> return ()
        _ -> assertFailure (show d)

  , testCase "if-then-else" $ do
      p <- parseOk "let x : Action = if a then Allow else Drop;"
      d <- singleDecl p
      case d of
        DLet _ _ (EIf (EVar "a") (EVar "Allow") (EVar "Drop")) -> return ()
        _ -> assertFailure (show d)

  , testCase "perform expression" $ do
      p <- parseOk "let x : Action = perform Log.emit(Info, \"msg\");"
      d <- singleDecl p
      case d of
        DLet _ _ (EPerform ["Log","emit"] [ELit (LString "Info"), ELit (LString "msg")]) -> return ()
        DLet _ _ (EPerform ["Log","emit"] _) -> return ()  -- arg parsing flexible
        _ -> assertFailure (show d)

  , testCase "do block" $ do
      p <- parseOk "let x : Action = do { y <- foo; y };"
      d <- singleDecl p
      case d of
        DLet _ _ (EDo [DSBind "y" _, DSExpr (EVar "y")]) -> return ()
        _ -> assertFailure (show d)

  , testCase "nested case" $ do
      p <- parseOk
            "let x : Action = case e of { \
            \  | a -> Allow; \
            \  | _ -> Drop; \
            \};"
      d <- singleDecl p
      case d of
        DLet _ _ (ECase (EVar "e") [Arm (PVar "a") Nothing _, Arm PWild Nothing _]) -> return ()
        _ -> assertFailure (show d)

  , testCase "lambda" $ do
      p <- parseOk "let x : Frame -> Action = \\frame -> Allow;"
      d <- singleDecl p
      case d of
        DLet _ _ (ELam "frame" (EVar "Allow")) -> return ()
        _ -> assertFailure (show d)

  , testCase "string concat" $ do
      p <- parseOk "let x : String = \"hello\" ++ \" world\";"
      d <- singleDecl p
      case d of
        DLet _ _ (EInfix OpConcat _ _) -> return ()
        _ -> assertFailure (show d)

  , testCase "negation" $ do
      p <- parseOk "let x : Bool = !flag;"
      d <- singleDecl p
      case d of
        DLet _ _ (ENot (EVar "flag")) -> return ()
        _ -> assertFailure (show d)

  , testCase "set literal" $ do
      p <- parseOk "let x : Set<Int> = { 22, 80, 443 };"
      d <- singleDecl p
      case d of
        DLet _ _ (ESet [ELit (LInt 22), ELit (LInt 80), ELit (LInt 443)]) -> return ()
        _ -> assertFailure (show d)
  ]

-- ─── Policy ──────────────────────────────────────────────────────────────────

policyTests :: TestTree
policyTests = testGroup "policy"
  [ testCase "minimal policy" $ do
      p <- parseOk
            "policy output : Frame \
            \  on { hook = Output, table = Filter, priority = Filter } \
            \  = { | _ -> Allow; };"
      d <- singleDecl p
      case d of
        DPolicy "output" _ (PolicyMeta HOutput TFilter (Priority 0)) [_] -> return ()
        _ -> assertFailure (show d)

  , testCase "NAT prerouting" $ do
      p <- parseOk
            "policy nat_pre : Frame \
            \  on { hook = Prerouting, table = NAT, priority = DstNat } \
            \  = { | _ -> Allow; };"
      d <- singleDecl p
      case d of
        DPolicy _ _ (PolicyMeta HPrerouting TNAT (Priority (-100))) _ -> return ()
        _ -> assertFailure (show d)

  , testCase "arm with guard" $ do
      p <- parseOk
            "policy input : Frame \
            \  on { hook = Input, table = Filter, priority = Filter } \
            \  = { \
            \    | _ if ct.state in { Established, Related } -> Allow; \
            \    | _ -> Drop; \
            \  };"
      d <- singleDecl p
      case d of
        DPolicy _ _ _ [Arm PWild (Just _) _, Arm PWild Nothing _] -> return ()
        _ -> assertFailure (show d)

  , testCase "Frame pattern with path" $ do
      p <- parseOk
            "policy forward : Frame \
            \  on { hook = Forward, table = Filter, priority = Filter } \
            \  = { \
            \    | Frame(iif in lan_zone -> wan, _) -> Allow; \
            \    | _ -> Drop; \
            \  };"
      d <- singleDecl p
      case d of
        DPolicy _ _ _ (Arm (PFrame (Just _) _) Nothing _ : _) -> return ()
        _ -> assertFailure (show d)

  , testCase "Frame pattern without Ether (layer stripping)" $ do
      p <- parseOk
            "policy input : Frame \
            \  on { hook = Input, table = Filter, priority = Filter } \
            \  = { \
            \    | Frame(_, IPv4(ip, TCP(tcp, _))) if tcp.dport == :22 -> Allow; \
            \    | _ -> Drop; \
            \  };"
      d <- singleDecl p
      case d of
        DPolicy _ _ _ (Arm (PFrame (Just _) (PCtor "IPv4" _)) _ _ : _) -> return ()
        _ -> assertFailure (show d)

  , testCase "policy arm calls rule" $ do
      p <- parseOk
            "policy forward : Frame \
            \  on { hook = Forward, table = Filter, priority = Filter } \
            \  = { \
            \    | frame -> blockOutboundWG(frame); \
            \  };"
      d <- singleDecl p
      case d of
        DPolicy _ _ _ [Arm (PVar "frame") Nothing (EApp (EVar "blockOutboundWG") _)] ->
          return ()
        _ -> assertFailure (show d)

  , testCase "Continue arm is parsed" $ do
      p <- parseOk
            "rule r : Frame -> Action = \
            \  \\frame -> case frame of { \
            \    | _ -> Continue; \
            \  };"
      d <- singleDecl p
      case d of
        DRule _ _ _ -> return ()
        _ -> assertFailure (show d)
  ]

-- ─── Rule ────────────────────────────────────────────────────────────────────

ruleTests :: TestTree
ruleTests = testGroup "rule"
  [ testCase "simple rule" $ do
      p <- parseOk
            "rule blockAll : Frame -> Action = \
            \  \\frame -> case frame of { | _ -> Drop; };"
      d <- singleDecl p
      case d of
        DRule "blockAll" _ (ELam "frame" (ECase _ _)) -> return ()
        _ -> assertFailure (show d)

  , testCase "rule with effects in type" $ do
      p <- parseOk
            "rule logged : Frame -> <Log> Action = \
            \  \\f -> case f of { | _ -> Allow; };"
      d <- singleDecl p
      case d of
        DRule "logged" (TFun _ (TEffect ["Log"] _)) _ -> return ()
        _ -> assertFailure (show d)

  , testCase "nested case in rule" $ do
      p <- parseOk
            "rule check : Frame -> <FlowMatch> Action = \
            \  \\frame -> \
            \    case frame of { \
            \      | Frame(_, IPv4(ip, UDP(udp, _))) -> \
            \          case perform FlowMatch.check(ip, wg) of { \
            \            | Matched -> Drop; \
            \            | _ -> Continue; \
            \          }; \
            \      | _ -> Continue; \
            \    };"
      d <- singleDecl p
      case d of
        DRule "check" _ (ELam _ (ECase _ _)) -> return ()
        _ -> assertFailure (show d)
  ]

-- ─── Config ──────────────────────────────────────────────────────────────────

configTests :: TestTree
configTests = testGroup "config"
  [ testCase "default table name" $ do
      p <- parseOk "interface wan : WAN {};"
      configTable (progConfig p) @?= "fwl"

  , testCase "custom table name" $ do
      p <- parseOk "config { table = \"myrules\"; } interface wan : WAN {};"
      configTable (progConfig p) @?= "myrules"
  ]

-- ─── Error cases ─────────────────────────────────────────────────────────────

errorTests :: TestTree
errorTests = testGroup "parse errors"
  [ testCase "missing semicolon" $
      parseFail "interface wan : WAN {}"

  , testCase "unknown hook" $
      parseFail
        "policy p : Frame \
        \  on { hook = Bogus, table = Filter, priority = Filter } \
        \  = { | _ -> Allow; };"

  , testCase "empty arm block with no arms is ok" $ do
      p <- parseOk
            "policy output : Frame \
            \  on { hook = Output, table = Filter, priority = Filter } \
            \  = {};"
      d <- singleDecl p
      case d of
        DPolicy _ _ _ [] -> return ()
        _ -> assertFailure (show d)

  , testCase "CIDR without prefix fails" $
      parseFail "interface lan : LAN { cidr4 = { 10.0.0.1 }; };"
  ]