open Ipaddr

type zone =
  | ZoneIpv4 of V4.Prefix.t
  | ZoneIpv6 of V6.Prefix.t
  | Zone of string
  | ZoneList of zone list
  | ZoneExclude of zone

type addrs = ZoneName of string | Ipv4 of V4.Prefix.t | Ipv6 of V6.Prefix.t

type l4_rule =
  | TcpRule of { sport : int list; dport : int list }
  | UdpRule of { sport : int list; dport : int list }
  | IcmpRule

type rule = { src : addrs list; dest : addrs list; l4 : l4_rule }
type config = { zones : (string * zone) list; rules : rule list }

let rec zone_of_json json =
  match json with
  | `String s -> (
      let zone = Zone s in
      let zone =
        match V6.Prefix.of_string s with Ok p -> ZoneIpv6 p | Error _ -> zone
      in
      match V4.Prefix.of_string s with Ok p -> ZoneIpv4 p | Error _ -> zone)
  | `Assoc [ ("exclude", e) ] -> ZoneExclude (zone_of_json e)
  | `List l -> ZoneList (List.map zone_of_json l)
  | _ -> assert false

let addrs_of_json json =
  let open Yojson.Basic.Util in
  let value = json |> to_string in
  let addrs = ZoneName value in
  let addrs =
    match V6.Prefix.of_string value with Ok p -> Ipv6 p | Error _ -> addrs
  in
  match V4.Prefix.of_string value with Ok p -> Ipv4 p | Error _ -> addrs

let addrs_list_of_json json =
  let open Yojson.Basic.Util in
  let elems = json |> to_list in
  List.map addrs_of_json elems

let to_list_force = function `List l -> l | _ -> []

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

let tcp_rule_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
  TcpRule { sport; dport }

let udp_rule_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
  UdpRule { sport; dport }

let l4_rule_of_json json =
  let open Yojson.Basic.Util in
  let proto = json |> member "proto" |> to_string in
  match proto with
  | "tcp" -> tcp_rule_of_json json
  | "udp" -> udp_rule_of_json json
  | "icmp" -> IcmpRule
  | _ -> assert false

let rule_of_json json =
  let open Yojson.Basic.Util in
  let src = addrs_list_of_json (json |> member "src") in
  let dest = addrs_list_of_json (json |> member "dest") in
  let l4 = l4_rule_of_json json in
  { src; dest; l4 }

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 config_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 |> List.map rule_of_json in
  { zones; rules }