(* Js_of_ocaml compiler
 * http://www.ocsigen.org/js_of_ocaml/
 * Copyright (C) 2019 Hugo Heuzard
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, with linking exception;
 * either version 2.1 of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 *)

(** Minimal version of ppx_optcomp
    It only support the following attribute
    {[
      [@if ocaml_version < (4,12,0)]
    ]}
    on module (Pstr_module),
    toplevel bindings (Pstr_value, Pstr_primitive)
    pattern in case (pc_lhs)
    and module in signature (Psig_module)
*)

open StdLabels
open Ppxlib.Parsetree

module Version : sig
  type t

  val of_list : int list -> t

  val compare : t -> t -> int

  val current : t
end = struct
  type t = int list

  let of_list l = l

  let split_char ~sep p =
    let len = String.length p in
    let rec split beg cur =
      if cur >= len
      then if cur - beg > 0 then [ String.sub p ~pos:beg ~len:(cur - beg) ] else []
      else if sep p.[cur]
      then String.sub p ~pos:beg ~len:(cur - beg) :: split (cur + 1) (cur + 1)
      else split beg (cur + 1)
    in
    split 0 0

  let split v =
    match
      split_char
        ~sep:(function
          | '+' | '-' | '~' -> true
          | _ -> false)
        v
    with
    | [] -> assert false
    | x :: _ ->
        List.map
          (split_char
             ~sep:(function
               | '.' -> true
               | _ -> false)
             x)
          ~f:int_of_string

  let current = split Sys.ocaml_version

  let compint (a : int) b = compare a b

  let rec compare v v' =
    match v, v' with
    | [ x ], [ y ] -> compint x y
    | [], [] -> 0
    | [], y :: _ -> compint 0 y
    | x :: _, [] -> compint x 0
    | x :: xs, y :: ys -> (
        match compint x y with
        | 0 -> compare xs ys
        | n -> n)
end

exception Invalid of Location.t

type op =
  | LE
  | GE
  | GT
  | LT
  | NEQ
  | EQ
  | AND
  | OR
  | NOT

type e =
  | Version of Version.t
  | Tuple of e list
  | Bool of bool
  | Int of int
  | String of string

let get_bin_op = function
  | { pexp_desc = Pexp_ident { txt = Lident str; _ }; pexp_loc = loc; _ } -> (
      match str with
      | "<=" -> LE
      | ">=" -> GE
      | ">" -> GT
      | "<" -> LT
      | "=" -> EQ
      | "<>" -> NEQ
      | "&&" -> AND
      | "||" -> OR
      | _ -> raise (Invalid loc))
  | { pexp_loc = loc; _ } -> raise (Invalid loc)

let get_un_op = function
  | { pexp_desc = Pexp_ident { txt = Lident str; _ }; pexp_loc = loc; _ } -> (
      match str with
      | "not" -> NOT
      | _ -> raise (Invalid loc))
  | { pexp_loc = loc; _ } -> raise (Invalid loc)

let version = function
  | Version v -> v
  | Tuple l ->
      Version.of_list
        (List.map l ~f:(function
          | Int i -> i
          | _ -> raise (Invalid Location.none)))
  | Bool _ | Int _ | String _ -> raise (Invalid Location.none)

let keep loc (attrs : attributes) =
  let ifs =
    List.filter attrs ~f:(function
      | { attr_name = { txt = "if"; _ }; _ } -> true
      | _ -> false)
  in
  match ifs with
  | [] -> true
  | _ -> (
      try
        let keep_one ({ attr_payload; attr_loc; _ } as attr) =
          Ppxlib.Attribute.mark_as_handled_manually attr;
          let e =
            match attr_payload with
            | PStr [ { pstr_desc = Pstr_eval (e, []); _ } ] -> e
            | _ -> raise (Invalid attr_loc)
          in
          let rec eval e =
            let open Ppxlib.Ast_pattern in
            let loc = e.pexp_loc in
            match
              (parse_res
                 (pexp_ident (lident (string "ocaml_version"))
                 >>| (fun () -> Version Version.current)
                 ||| (pexp_ident (lident (string "ast_version"))
                     >>| fun () -> Int Ppxlib.Selected_ast.version)
                 ||| (pexp_construct (lident (string "true")) drop >>| fun () -> Bool true)
                 ||| (pexp_construct (lident (string "false")) drop
                     >>| fun () -> Bool false)
                 ||| (pexp_constant (pconst_integer __ none)
                     >>| fun () d -> Int (int_of_string d))
                 ||| (pexp_tuple __ >>| fun () l -> Tuple (List.map l ~f:eval))
                 ||| (pexp_apply __ __
                     >>| fun () op l ->
                     match l with
                     | [ (Nolabel, a); (Nolabel, b) ] -> (
                         let op = get_bin_op op in
                         let a = eval a in
                         let b = eval b in
                         match op with
                         | LE | GE | LT | GT | NEQ | EQ ->
                             let comp =
                               match a, b with
                               | Version _, _ | _, Version _ ->
                                   Version.compare (version a) (version b)
                               | Int a, Int b -> compare a b
                               | _ -> raise (Invalid loc)
                             in
                             let op =
                               match op with
                               | LE -> ( <= )
                               | GE -> ( >= )
                               | LT -> ( < )
                               | GT -> ( > )
                               | EQ -> ( = )
                               | NEQ -> ( <> )
                               | _ -> assert false
                             in
                             Bool (op comp 0)
                         | AND -> (
                             match a, b with
                             | Bool a, Bool b -> Bool (a && b)
                             | _ -> raise (Invalid loc))
                         | OR -> (
                             match a, b with
                             | Bool a, Bool b -> Bool (a || b)
                             | _ -> raise (Invalid loc))
                         | NOT -> raise (Invalid loc))
                     | [ (Nolabel, a) ] -> (
                         let op = get_un_op op in
                         let a = eval a in
                         match op, a with
                         | NOT, Bool b -> Bool (not b)
                         | NOT, _ -> raise (Invalid loc)
                         | _ -> raise (Invalid loc))
                     | _ -> raise (Invalid loc))))
                loc
                e
                ()
            with
            | Ok res -> res
            | Error _ -> raise (Invalid loc)
          in
          match eval e with
          | Bool b -> b
          | Int _ | String _ | Tuple _ | Version _ -> raise (Invalid loc)
        in
        let keep = List.for_all ~f:keep_one ifs in
        if false
        then
          if not keep
          then
            Printf.eprintf
              "dropping %s:%d\n%!"
              loc.Location.loc_start.pos_fname
              loc.Location.loc_start.pos_lnum
          else
            Printf.eprintf
              "keep %s:%d\n%!"
              loc.Location.loc_start.pos_fname
              loc.Location.loc_start.pos_lnum;
        keep
      with Invalid loc -> Location.raise_errorf ~loc "Invalid attribute format")

let filter_map ~f l =
  let l =
    List.fold_left
      ~f:(fun acc x ->
        match f x with
        | Some x -> x :: acc
        | None -> acc)
      ~init:[]
      l
  in
  List.rev l

let rec filter_pattern = function
  | { ppat_desc = Ppat_or (p1, p2); _ } as p -> (
      match filter_pattern p1, filter_pattern p2 with
      | None, None -> None
      | Some p1, None -> Some p1
      | None, Some p2 -> Some p2
      | Some p1, Some p2 -> Some { p with ppat_desc = Ppat_or (p1, p2) })
  | { ppat_attributes; ppat_loc; _ } as p ->
      if keep ppat_loc ppat_attributes then Some p else None

let drop_attr =
  let attr =
    { attr_name = Location.mknoloc "ppx_optcomp_light.dropped"
    ; attr_loc = Location.none
    ; attr_payload = PStr []
    }
  in
  Ppxlib.Attribute.mark_as_handled_manually attr;
  attr

let drop_str loc = { pstr_desc = Pstr_attribute drop_attr; pstr_loc = loc }

let drop_sig loc = { psig_desc = Psig_attribute drop_attr; psig_loc = loc }

let traverse =
  object
    inherit Ppxlib.Ast_traverse.map as super

    method! structure_item item =
      let item = super#structure_item item in
      match item.pstr_desc with
      | Pstr_module { pmb_attributes; pmb_loc; _ } ->
          if keep pmb_loc pmb_attributes then item else drop_str pmb_loc
      | Pstr_primitive { pval_attributes; pval_loc; _ } ->
          if keep pval_loc pval_attributes then item else drop_str pval_loc
      | Pstr_value (r, l) -> (
          let l =
            filter_map l ~f:(fun b ->
                if keep b.pvb_loc b.pvb_attributes then Some b else None)
          in
          match l with
          | [] -> drop_str Location.none
          | _ -> { item with pstr_desc = Pstr_value (r, l) })
      | _ -> item

    method! cases cases =
      let cases =
        filter_map cases ~f:(fun case ->
            match filter_pattern case.pc_lhs with
            | None -> None
            | Some pattern -> Some { case with pc_lhs = pattern })
      in
      super#cases cases

    method! signature_item item =
      let item = super#signature_item item in
      match item.psig_desc with
      | Psig_module { pmd_attributes; pmd_loc; _ } ->
          if keep pmd_loc pmd_attributes then item else drop_sig pmd_loc
      | _ -> item
  end

let () =
  Ppxlib.Driver.register_transformation
    ~impl:traverse#structure
    ~intf:traverse#signature
    "ppx_optcomp_light"
