module CheckTests (tests) where

import Test.Tasty
import Test.Tasty.HUnit

import FWL.Check
import FWL.Util

tests :: TestTree
tests = testGroup "Check"
  [ undefinedNameTests
  , duplicateTests
  , policyTerminationTests
  , patternCycleTests
  , cleanProgramTests
  ]

-- ─── Helper ──────────────────────────────────────────────────────────────────

checkSrc :: String -> IO [CheckError]
checkSrc src = do
  p <- parseOk src
  return (checkProgram p)

assertNoErrors :: String -> IO ()
assertNoErrors src = do
  errs <- checkSrc src
  case errs of
    [] -> return ()
    _  -> assertFailure ("Unexpected errors: " ++ show errs)

assertHasError :: (CheckError -> Bool) -> String -> IO ()
assertHasError p src = do
  errs <- checkSrc src
  if any p errs
    then return ()
    else assertFailure ("Expected error not found. Got: " ++ show errs)

isUndefined :: String -> CheckError -> Bool
isUndefined n (UndefinedName _ m) = m == n
isUndefined _ _                   = False

isDuplicate :: String -> CheckError -> Bool
isDuplicate n (DuplicateDecl _ m) = m == n
isDuplicate _ _                   = False

isNoContinue :: String -> CheckError -> Bool
isNoContinue n (PolicyNoContinue m) = m == n
isNoContinue _ _                    = False

isCycle :: CheckError -> Bool
isCycle (PatternCycle _) = True
isCycle _                = False

-- ─── Undefined name tests ────────────────────────────────────────────────────

undefinedNameTests :: TestTree
undefinedNameTests = testGroup "undefined names"
  [ testCase "zone references unknown interface" $
      assertHasError (isUndefined "ghost")
        "zone bad_zone = { lan, ghost };"

  , testCase "zone references known interface — no error" $
      assertNoErrors
        "interface lan : LAN {}; \
        \zone good = { lan };"

  , testCase "pattern references undefined named pattern" $
      assertHasError (isUndefined "Undefined")
        "pattern Bad : Frame = Frame(_, IPv4(ip, Undefined));"

  , testCase "pattern references known named pattern — no error" $
      assertNoErrors
        "pattern WGInit : (UDPHeader,Bytes) = (udp { length = 156 }, [0x01 _*]); \
        \pattern Compound : Frame = Frame(_, IPv4(ip, WGInit));"

  , testCase "flow references undefined pattern" $
      assertHasError (isUndefined "Ghost")
        "flow Bad : FlowPattern = Ghost;"

  , testCase "flow references known pattern — no error" $
      assertNoErrors
        "pattern P : T = udp { length = 1 }; \
        \flow F : FlowPattern = P;"

  , testCase "policy guard references undeclared zone" $
      -- 'unknown_zone' not declared; check should flag it
      assertHasError (isUndefined "unknown_zone")
        "policy fwd : Frame \
        \  on { hook = Forward, table = Filter, priority = Filter } \
        \  = { | Frame(iif in unknown_zone -> wan, _) -> Allow; \
        \      | _ -> Drop; \
        \  };"

  , testCase "policy references known zone — no error" $
      assertNoErrors
        "interface lan : LAN {}; \
        \zone trusted = { lan }; \
        \policy fwd : Frame \
        \  on { hook = Forward, table = Filter, priority = Filter } \
        \  = { | Frame(iif in trusted -> wan, _) -> Allow; \
        \      | _ -> Drop; \
        \  };"
  ]

-- ─── Duplicate declaration tests ─────────────────────────────────────────────

duplicateTests :: TestTree
duplicateTests = testGroup "duplicates"
  [ testCase "duplicate interface" $
      assertHasError (isDuplicate "lan")
        "interface lan : LAN {}; \
        \interface lan : WAN {};"

  , testCase "duplicate zone" $
      assertHasError (isDuplicate "z")
        "zone z = { a }; \
        \zone z = { b };"

  , testCase "duplicate pattern" $
      assertHasError (isDuplicate "P")
        "pattern P : T = udp { length = 1 }; \
        \pattern P : T = udp { length = 2 };"

  , testCase "duplicate policy" $
      assertHasError (isDuplicate "input")
        "policy input : Frame \
        \  on { hook = Input, table = Filter, priority = Filter } \
        \  = { | _ -> Allow; }; \
        \policy input : Frame \
        \  on { hook = Input, table = Filter, priority = Filter } \
        \  = { | _ -> Drop; };"

  , testCase "distinct names — no error" $
      assertNoErrors
        "interface lan : LAN {}; \
        \interface wan : WAN { dynamic; }; \
        \zone z = { lan };"
  ]

-- ─── Policy termination tests ────────────────────────────────────────────────

policyTerminationTests :: TestTree
policyTerminationTests = testGroup "policy termination"
  [ testCase "last arm is Continue — error" $
      assertHasError (isNoContinue "bad_policy")
        "policy bad_policy : Frame \
        \  on { hook = Input, table = Filter, priority = Filter } \
        \  = { | _ -> Continue; };"

  , testCase "last arm is Drop — ok" $
      assertNoErrors
        "policy good : Frame \
        \  on { hook = Input, table = Filter, priority = Filter } \
        \  = { | _ if ct.state in { Established } -> Allow; \
        \      | _ -> Drop; \
        \  };"

  , testCase "last arm is Allow — ok" $
      assertNoErrors
        "policy output : Frame \
        \  on { hook = Output, table = Filter, priority = Filter } \
        \  = { | _ -> Allow; };"

  , testCase "Continue in non-last arm is fine" $
      assertNoErrors
        "rule r : Frame -> Action = \
        \  \\f -> case f of { \
        \    | Frame(_, IPv4(ip, _)) -> Continue; \
        \    | _ -> Drop; \
        \  };"

  , testCase "empty policy body — error" $
      assertHasError (isNoContinue "empty")
        "policy empty : Frame \
        \  on { hook = Output, table = Filter, priority = Filter } \
        \  = {};"
  ]

-- ─── Pattern cycle tests ─────────────────────────────────────────────────────

patternCycleTests :: TestTree
patternCycleTests = testGroup "pattern cycles"
  [ testCase "direct self-reference — cycle error" $
      assertHasError isCycle
        "pattern Loop : T = Frame(_, Loop);"

  , testCase "mutual cycle — cycle error" $
      assertHasError isCycle
        "pattern A : T = Frame(_, B); \
        \pattern B : T = Frame(_, A);"

  , testCase "linear chain — no cycle" $
      assertNoErrors
        "pattern Base : T = udp { length = 1 }; \
        \pattern Mid  : T = Frame(_, Base); \
        \pattern Top  : T = Frame(_, Mid);"
  ]

-- ─── Clean full programs ──────────────────────────────────────────────────────

cleanProgramTests :: TestTree
cleanProgramTests = testGroup "clean programs"
  [ testCase "minimal router skeleton" $
      assertNoErrors
        "interface wan : WAN { dynamic; }; \
        \interface lan : LAN { cidr4 = { 10.17.1.0/24 }; }; \
        \interface wg0 : WireGuard {}; \
        \zone lan_zone = { lan, wg0 }; \
        \policy input : Frame \
        \  on { hook = Input, table = Filter, priority = Filter } \
        \  = { | _ if ct.state in { Established, Related } -> Allow; \
        \      | _ -> Drop; \
        \  }; \
        \policy output : Frame \
        \  on { hook = Output, table = Filter, priority = Filter } \
        \  = { | _ -> Allow; };"

  , testCase "pattern and flow declarations" $
      assertNoErrors
        "pattern WGInit : (UDPHeader,Bytes) = (udp { length = 156 }, [0x01 _*]); \
        \pattern WGResp : (UDPHeader,Bytes) = (udp { length = 100 }, [0x02 _*]); \
        \flow WGHandshake : FlowPattern = WGInit . WGResp within 5s;"
  ]