open Utils
open Ipaddr

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

module Payload = struct
  module Udp = struct
    type _ t = Sport : int t | Dport : int t

    let to_string : type a. a t -> string = function
      | Sport -> "sport"
      | Dport -> "dport"
  end

  module Tcp = struct
    type _ t = Sport : int t | Dport : int t

    let to_string : type a. a t -> string = function
      | Sport -> "sport"
      | Dport -> "dport"
  end

  module Ipv4 = struct
    type _ t = Saddr : V4.Prefix.t t | Daddr : V4.Prefix.t t

    let to_string : type a. a t -> string = function
      | Saddr -> "saddr"
      | Daddr -> "daddr"
  end

  module Ipv6 = struct
    type _ t = Saddr : V6.Prefix.t t | Daddr : V6.Prefix.t t

    let to_string : type a. a t -> string = function
      | Saddr -> "saddr"
      | Daddr -> "daddr"
  end

  type _ t =
    | Udp : 'a Udp.t -> 'a t
    | Tcp : 'a Tcp.t -> 'a t
    | Ipv4 : 'a Ipv4.t -> 'a t
    | Ipv6 : 'a Ipv6.t -> 'a t

  let to_json (type a) (payload : a t) =
    let protocol, field =
      match payload with
      | Udp udp -> ("udp", Udp.to_string udp)
      | Tcp tcp -> ("tcp", Tcp.to_string tcp)
      | Ipv4 ipv4 -> ("ip", Ipv4.to_string ipv4)
      | Ipv6 ipv6 -> ("ip6", Ipv6.to_string ipv6)
    in
    assoc_one "payload"
      (`Assoc [ ("protocol", `String protocol); ("field", `String field) ])
end

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

  let rec to_json : type a. a t -> Yojson.Basic.t = function
    | String str -> `String str
    | Number num -> `Int num
    | Boolean bool -> `Bool bool
    | Ipv4 ipv4 ->
        let network = V4.(to_string (Prefix.network ipv4)) in
        let len = V4.Prefix.bits ipv4 in
        assoc_one "prefix"
          (`Assoc [ ("addr", `String network); ("len", `Int len) ])
    | Ipv6 ipv6 ->
        let network = V6.(to_string (Prefix.network ipv6)) in
        let len = V6.Prefix.bits ipv6 in
        assoc_one "prefix"
          (`Assoc [ ("addr", `String network); ("len", `Int len) ])
    | List list -> `List (List.map to_json list)
    | Set set -> assoc_one "set" (`List (List.map to_json set))
    | Range (a, b) -> assoc_one "range" (`List [ to_json a; to_json b ])
    | Payload payload -> Payload.to_json payload
end

module Match = struct
  type (_, _) t = Eq : ('a, 'a) t | NotEq : ('a, 'a) t

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

  let to_json (type a b) (op : (a, b) t) = `String (to_string op)
end

module Counter = struct
  type t =
    | NamedCounter of string
    | AnonCounter of { packets : int; bytes : int }

  let to_json = function
    | NamedCounter n -> assoc_one "counter" (`String n)
    | AnonCounter { packets; bytes } ->
        assoc_one "counter"
          (`Assoc [ ("packets", `Int packets); ("bytes", `Int bytes) ])
end

module Verdict = struct
  type t = Accept | Drop | Continue | Return | Jump of string | Goto of string

  let to_json = 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))
end

module Stmt = struct
  type _ t =
    | Match : (('a, 'b) Match.t * 'a Expr.t * 'b Expr.t) -> unit t
    | Counter : Counter.t -> unit t
    | Verdict : Verdict.t -> unit t
    | NoTrack : unit t
    | Log : { prefix : string option; group : int option } -> unit t

  let to_json : type a. a t -> Yojson.Basic.t = function
    | Match (op, left, right) ->
        assoc_one "match"
          (`Assoc
            [
              ("left", Expr.to_json left);
              ("right", Expr.to_json right);
              ("op", Match.to_json op);
            ])
    | Counter counter -> Counter.to_json counter
    | Verdict verdict -> Verdict.to_json verdict
    | 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))
end

module Family = struct
  type t = Ipv4 | Ipv6 | Inet

  let to_string = function Ipv4 -> "ip" | Ipv6 -> "ip6" | Inet -> "inet"
  let to_json f = `String (to_string f)
end

module Table = struct
  type t = { family : Family.t; name : string }

  let to_json { family; name } =
    assoc_one "table"
      (`Assoc [ ("family", Family.to_json family); ("name", `String name) ])
end

module Chain = struct
  type t = { family : Family.t; table : string; name : string }

  let to_json { family; table; name } =
    assoc_one "chain"
      (`Assoc
        [
          ("family", Family.to_json family);
          ("table", `String table);
          ("name", `String name);
        ])
end

module Rule = struct
  type t = {
    family : Family.t;
    table : string;
    chain : string;
    expr : unit Stmt.t list;
  }

  let to_json { family; table; chain; expr } =
    assoc_one "rule"
      (`Assoc
        [
          ("family", Family.to_json family);
          ("table", `String table);
          ("chain", `String chain);
          ("expr", `List (List.map Stmt.to_json expr));
        ])
end

module Command = struct
  type t =
    | FlushRuleset
    | AddTable of Table.t
    | FlushTable of Table.t
    | AddChain of Chain.t
    | FlushChain of Chain.t
    | AddRule of Rule.t

  let to_json = function
    | FlushRuleset -> assoc_one "flush" (assoc_one "ruleset" `Null)
    | AddTable table -> assoc_one "add" (Table.to_json table)
    | FlushTable table -> assoc_one "flush" (Table.to_json table)
    | AddChain chain -> assoc_one "add" (Chain.to_json chain)
    | FlushChain chain -> assoc_one "flush" (Chain.to_json chain)
    | AddRule rule -> assoc_one "add" (Rule.to_json rule)
end

let to_json commands =
  assoc_one "nftables" (`List (List.map Command.to_json commands))