open Utils
open Ipaddr

type _ udp =
  | UdpDport : int udp
  | UdpSport : int udp

type _ tcp =
  | TcpDport : int tcp
  | TcpSport : int tcp

type _ payload =
  | Udp : 'a udp -> 'a payload
  | Tcp : 'a tcp -> 'a payload

type _ expr =
  | String : string -> string expr
  | Number : int -> int expr
  | Boolean : bool -> int expr
  | Ipv4 : (V4.t * V4.Prefix.t) -> (V4.t * V4.Prefix.t) expr
  | Ipv6 : (V6.t * V6.Prefix.t) -> (V6.t * V6.Prefix.t) expr
  | List : 'a expr list -> 'a expr
  | Set : 'a expr list -> 'a expr
  | Range : 'a expr * 'a expr -> 'a expr
  | Payload : 'a payload -> 'a expr

type (_, _) match_op =
  | Eq : ('a, 'a) match_op
  | NotEq : ('a, 'a) match_op

type counter =
  | Named of string
  | Anon of { packets: int; bytes: int }

type verdict =
  | Accept
  | Drop
  | Continue
  | Return
  | Jump of string
  | Goto of string

type _ stmt =
  | Match : {
      left: 'a expr;
      right: 'b expr; 
      op: ('a, 'b) match_op
    } -> unit stmt
  | Counter : counter -> unit stmt
  | Verdict : verdict -> unit stmt
  | NoTrack : unit stmt
  | Log : {
      prefix: string option;
      group: int option
    } -> unit stmt

type family = Ip6 | Ip4 | Inet

type table = { family: family; table_name: string }

type chain = { family: family; table: string; chain_name: string }

type rule = {
  family: family;
  table: string;
  chain: string;
  expr: unit stmt list
}

type add_object =
  | AddTable of table
  | AddChain of chain
  | AddRule of rule

type flush_object =
  | FlushRuleset
  | FlushTable of table
  | FlushChain of chain

type command =
  | Add of add_object
  | Flush of flush_object

let string_of_udp : type a. a udp -> string = function
  | UdpSport -> "sport"
  | UdpDport -> "dport"

let string_of_tcp : type a. a tcp -> string = function
  | TcpSport -> "sport"
  | TcpDport -> "dport"

let assoc_one key value = `Assoc [(key, value)]

let json_of_payload (type a) (payload : a payload) =
  let (protocol, field) = match payload with
    | Udp udp -> "udp", string_of_udp udp
    | Tcp tcp -> "tcp", string_of_tcp tcp
  in
  assoc_one "payload" (`Assoc [
      ("protocol", `String protocol);
      ("field", `String field);
    ])

let rec json_of_expr : type a. a expr -> Yojson.Basic.t = function
  | String s -> `String s
  | Number n -> `Int n
  | Boolean b -> `Bool b
  | Ipv4 (a, p) -> assoc_one "prefix" (`Assoc [
      ("addr", `String (V4.to_string a));
      ("len", `Int (V4.Prefix.bits p))
    ])
  | Ipv6 (a, p) -> assoc_one "prefix" (`Assoc [
      ("addr", `String (V6.to_string a));
      ("len", `Int (V6.Prefix.bits p))
    ])
  | List l -> `List (List.map json_of_expr l)
  | Set s -> assoc_one "set" (`List (List.map json_of_expr s))
  | Range (a, b) ->
    assoc_one "range" (`List [json_of_expr a ; json_of_expr b])
  | Payload p -> json_of_payload p

let string_of_match_op : type a b. (a, b) match_op -> string = function
  | Eq -> "=="
  | NotEq -> "!="

let json_of_counter = function
  | Named n -> assoc_one "counter" (`String n)
  | Anon { packets; bytes } ->
    assoc_one "counter" (`Assoc [
        ("packets", `Int packets);
        ("bytes", `Int bytes)
      ])

let json_of_verdict = function
  | Accept -> assoc_one "accept" `Null
  | Drop -> assoc_one "drop" `Null
  | Continue -> assoc_one "continue" `Null
  | Return -> assoc_one "return" `Null
  | Jump s -> assoc_one "jump" (assoc_one "target" (`String s))
  | Goto s -> assoc_one "goto" (assoc_one "target" (`String s))

let json_of_stmt : type a. a stmt -> Yojson.Basic.t = function
  | Match { left; right; op } ->
    assoc_one "match" (`Assoc [
        ("left", json_of_expr left);
        ("right", json_of_expr right);
        ("op", `String (string_of_match_op op))
      ])
  | Counter c -> json_of_counter c
  | Verdict v -> json_of_verdict v
  | NoTrack -> assoc_one "notrack" `Null
  | Log { prefix; group } ->
    let elems = [
      Option.map (fun p -> ("prefix", `String p)) prefix;
      Option.map (fun g -> ("group", `Int g)) group
    ] in
    assoc_one "log" (`Assoc (deoptionalise elems))

let string_of_family = function
  | Ip6 -> "ip6"
  | Ip4 -> "ip4"
  | Inet -> "inet"

let json_of_table { family; table_name } =
  assoc_one "table" (`Assoc [
      ("family", `String (string_of_family family));
      ("name", `String table_name)
    ])

let json_of_chain { family; table; chain_name } =
  assoc_one "chain" (`Assoc [
      ("family", `String (string_of_family family));
      ("table", `String table);
      ("name", `String chain_name)
    ])

let json_of_rule { family; table; chain; expr } =
  assoc_one "rule" (`Assoc [
      ("family", `String (string_of_family family));
      ("table", `String table);
      ("chain", `String chain);
      ("expr", `List (List.map json_of_stmt expr))
    ])

let json_of_add_object = function
  | AddTable t -> json_of_table t
  | AddChain c -> json_of_chain c
  | AddRule r -> json_of_rule r

let json_of_flush_object = function
  | FlushRuleset -> assoc_one "ruleset" `Null
  | FlushTable t -> json_of_table t
  | FlushChain c -> json_of_chain c

let json_of_command = function
  | Add a -> assoc_one "add" (json_of_add_object a)
  | Flush a -> assoc_one "flush" (json_of_flush_object a)

let json_of_nftables n =
  assoc_one "nftables" (`List (List.map json_of_command n))