open Ipaddr
open Utils

module Zone = struct
  type t =
    | Ipv4 of V4.Prefix.t
    | Ipv6 of V6.Prefix.t
    | Name of string
    | List of t list
    | Not of t

  let name str = Name str
  let ipv4 prefix = Ipv4 prefix
  let ipv6 prefix = Ipv6 prefix

  let rec of_json = function
    | `String str ->
        str |> V6.Prefix.of_string |> Result.map ipv6
        &?> (str |> V4.Prefix.of_string |> Result.map ipv4)
        &> (str |> name)
    | `Assoc [ ("not", json) ] -> Not (of_json json)
    | `List list -> List (List.map of_json list)
    | _ -> failwith "invalid zone definition"
end

module Addrs = struct
  type t = Name of string | Ipv4 of V4.Prefix.t | Ipv6 of V6.Prefix.t

  let name str = Name str
  let ipv4 prefix = Ipv4 prefix
  let ipv6 prefix = Ipv6 prefix

  let of_json json =
    let open Yojson.Basic.Util in
    let str = json |> to_string in
    str |> V6.Prefix.of_string |> Result.map ipv6
    &?> (str |> V4.Prefix.of_string |> Result.map ipv4)
    &> (str |> name)
end

let to_list_loose = function `List list -> list | _ -> []

let to_int_list json =
  let open Yojson.Basic.Util in
  json |> to_list_loose |> List.map to_int

module PayloadRule = struct
  module Tcp = struct
    type t = { sport : int list; dport : int list }

    let of_json json =
      let open Yojson.Basic.Util in
      let sport = json |> member "sport" |> to_int_list in
      let dport = json |> member "dport" |> to_int_list in
      { sport; dport }
  end

  module Udp = struct
    type t = { sport : int list; dport : int list }

    let of_json json =
      let open Yojson.Basic.Util in
      let sport = json |> member "sport" |> to_int_list in
      let dport = json |> member "dport" |> to_int_list in
      { sport; dport }
  end

  type t = Tcp of Tcp.t | Udp of Udp.t | Icmp

  let of_json json =
    let open Yojson.Basic.Util in
    match json |> member "proto" |> to_string with
    | "tcp" -> Tcp (Tcp.of_json json)
    | "udp" -> Udp (Udp.of_json json)
    | "icmp" -> Icmp
    | proto -> failwith ("invalid protocol " ^ proto)
end

let to_addr_list json = json |> to_list_loose |> List.map Addrs.of_json

module Rule = struct
  type t = { src : Addrs.t list; dest : Addrs.t list; payload : PayloadRule.t }

  let of_json json =
    let open Yojson.Basic.Util in
    let src = json |> member "src" |> to_addr_list in
    let dest = json |> member "dest" |> to_addr_list in
    let payload = PayloadRule.of_json json in
    { src; dest; payload }
end

type t = { zones : (string * Zone.t) list; rules : Rule.t list }

let zones_of_json json =
  let open Yojson.Basic.Util in
  json |> to_assoc |> List.map (fun (n, z) -> (n, Zone.of_json z))

let of_json json =
  let open Yojson.Basic.Util in
  let zones = json |> member "zones" |> zones_of_json in
  let rules =
    json |> member "rules" |> to_list_loose |> List.map Rule.of_json
  in
  { zones; rules }