exception Unexpected_null of string * string * string

module Idmap = Map.Make(struct type t = string let compare = compare end)
type env = string Idmap.t

type 'a order = Asc of 'a * string | Desc of 'a * string

type ('a, 'aa) conversion_functions = {
  value_to_sql : 'a -> string;
  value_of_sql : 'aa -> bool -> string -> 'a;
}

type ('a, 'aa) relation =
    Table of string * ('a, 'aa) conversion_functions * string list
  (* | Project of exists 'b 'bb. ('b, 'bb) relation * 'a -> 'b *)
  | Project of ('a, 'aa) packed_projection * string list
  | Select of ('a, 'aa) relation * ('a, 'aa) criterion
  (* | Rename of exists 'b 'bb. ('b, 'bb) relation * 'a -> 'b *)
  | Rename of ('a, 'aa) packed_projection * (string * string) list
  (* | Product of exists 'b 'bb 'c 'cc 'b1 'c1.
        ('b, 'bb) relation * ('b1 -> 'a) * ('b1 -> b) *
        ('c, 'cc) relation * ('c1 -> 'a) * ('c1 -> 'c) *)
  (* the sum of 'b and 'c is a subtype of 'a
   * 'b is a subtype of 'b1 and 'c is a subtype of 'c1
   * 'a is a subtype of 'b1 and 'c1
   * ditto for 'aa, 'bb and 'cc *)
  | Product of ('a, 'aa) packed_product * string list * string list
  | Order of ('a, 'aa) relation * 'aa order list
  | Limit of ('a, 'aa) relation * int
  | Offset of ('a, 'aa) relation * int

and ('a, 'aa, 'b, 'bb) projection_data = {
  proj_conv : ('a, 'aa) conversion_functions;
  proj_rel : ('b, 'bb) relation;
  proj_f : ('a -> 'b);
  proj_to_sql : env -> ('b, 'bb) relation -> env * string
}

and ('a, 'aa, 'z) projection_scope = {
  bind_projection : 'b 'b1.  ('a, 'aa, 'b, 'b1) projection_data -> 'z
}

and ('a, 'aa) packed_projection = {
  open_projection : 'z. ('a, 'aa, 'z) projection_scope -> 'z
}

and ('a, 'aa) packed_product = {
  open_product : 'z. ('a, 'aa, 'z) product_scope -> 'z
}

and ('a, 'aa, 'z) product_scope = {
  bind_product : 'b 'c 'b1 'c1 'bb 'cc.
    ('a, 'aa, 'b, 'c, 'b1, 'c1, 'bb, 'cc) product_data -> 'z
}

and ('a, 'aa, 'b, 'c, 'b1, 'c1, 'bb, 'cc) product_data = {
  prod_conv : ('a, 'aa) conversion_functions;
  prod_rel1 : ('b, 'bb) relation;
  prod_rel2 : ('c, 'cc) relation;
  prod_to_sql1 : env -> ('b, 'bb) relation -> env * string;
  prod_to_sql2 : env -> ('c, 'cc) relation -> env * string;
  prod_f : ('b1 -> 'a) * ('b1 -> 'b) * ('c1 -> 'a) * ('c1 -> 'c)
}

and operator = Eq | Neq | Lt | Gt | Le | Ge | Like

and ('a, 'aa) criterion =
  Sel_and of ('a, 'aa) criterion list
  | Sel_or of ('a, 'aa) criterion list
  | Sel_binop of operator * 'a * 'aa * string
  | Sel_binop' of operator * 'aa * 'aa * string * string * ('a * 'a -> bool)
  (* ... *)

let string_of_operator = function
    Eq -> "="
  | Neq -> "<>"
  | Lt -> "<"
  | Gt -> ">"
  | Le -> "<="
  | Ge -> ">="
  | Like -> "LIKE"

let with_packed_projection p e = p.open_projection e

let with_packed_product p e = p.open_product e

let select_fields fields =
  "SELECT " ^ (String.concat ", " (List.map (fun x -> "\"" ^ x ^ "\"") fields))

let select_fields_with_prefix l =
  "SELECT " ^ (String.concat ", "
                 (List.map (fun (pre, x) -> pre ^ "." ^ "\"" ^ x ^ "\"") l))

let projection_to_sql_subquery t =
  with_packed_projection t
    { bind_projection = fun proj -> proj.proj_to_sql Idmap.empty proj.proj_rel }

let merge_environments old newer = Idmap.fold Idmap.add newer old

let product_to_sql_subquery id1 id2 t =
  with_packed_product t
    { bind_product = fun prod ->
        let env1, t1 = prod.prod_to_sql1 Idmap.empty prod.prod_rel1 in
        let env2, t2 = prod.prod_to_sql2 Idmap.empty prod.prod_rel2 in
        let env = merge_environments env1 env2 in
          (env, t1 ^ " AS " ^ id1 ^ ", " ^ t2 ^ " AS " ^ id2) }

let new_sql_identifier =
  let id = ref 0 in
    fun () -> incr id; Printf.sprintf "t%d" !id

let rec conversion_functions = function
    Table (_, f, _) -> f
  | Project (t, _) | Rename (t, _) ->
      with_packed_projection t { bind_projection = fun proj -> proj.proj_conv }
  | Select (rel, _) -> conversion_functions rel
  | Product (t, _, _) ->
      with_packed_product t { bind_product = fun prod -> prod.prod_conv }
  | Order (rel, _) | Limit (rel, _) | Offset (rel, _) -> conversion_functions rel

let get_full_field_name env field =
  let esc_field = "\"" ^ field ^ "\"" in
    try
      (Idmap.find field env) ^ "." ^ esc_field
    with Not_found -> esc_field

let criterion_to_sql rel env criterion =
  let rec aux = function
      Sel_and l -> String.concat " AND "
                     (List.map (fun txt -> "(" ^ txt ^ ")") (List.map aux l))
    | Sel_or l -> String.concat " OR "
                    (List.map (fun txt -> "(" ^ txt ^ ")") (List.map aux l))
    | Sel_binop (op, x, ftype, field) ->
        let to_s = (conversion_functions rel).value_to_sql in
          get_full_field_name env field ^ " " ^ string_of_operator op ^ " " ^ to_s x
    | Sel_binop' (op, _, _, field1, field2, _) ->
        get_full_field_name env field1 ^ " " ^ string_of_operator op ^ " " ^
        get_full_field_name env field2

  in aux criterion

let rec to_sql rel =
  let rel = simplify_rel_tree rel in
  let _, text = to_sql_aux Idmap.empty rel in text

and criteria_equal c1 c2 = match (c1, c2) with
    Sel_and a, Sel_and b | Sel_or a, Sel_or b ->
      List.length a = List.length b &&
      List.fold_left2 (fun s x y -> s && criteria_equal x y) true a b
  | Sel_binop (op1, v1, f1, n1), Sel_binop (op2, v2, f2, n2) ->
      op1 = op2 && v1 = v2 && f1 = f2 && n1 = n2
  | Sel_binop' (opa, fa1, fa2, na1, na2, _),
    Sel_binop' (opb, fb1, fb2, nb1, nb2, _) ->
      opa = opb && fa1 = fb1 && fa2 = fb2 && na1 = nb1 && na2 = nb2
  | _, _ -> false

(* we use physical equality to know if the children of an operator have been
 * simplified, making further simplifications possible  *)
and maybe_simplify f ~child old =
  let child' = simplify_rel_tree child in
    if child' == child then old
    else simplify_rel_tree (f child')

(* it is imperative that the very node we're simplifying (not a copy) be
 * returned when no simplification has taken place so that simplify_rel_tree
 * terminates *)
and simplify_rel_tree = function
    Order (Limit(r, l), o) -> simplify_rel_tree (Limit (Order(r, o), l))
  | Limit (Limit (r, l1), l2) -> simplify_rel_tree (Limit (r, min l1 l2))
  | Order(Offset(r, o1), o2) -> simplify_rel_tree (Offset(Order(r, o2), o1))
  | Offset(Limit(r, l1), o) -> simplify_rel_tree (Limit(Offset(r, o), l1))
  | Offset(Offset(r, o1), o2) -> simplify_rel_tree (Offset(r, o1+o2))
  | Order (Order (r, _), o) -> simplify_rel_tree (Order (r, o))
  | Select (Select (r, c1), c2) ->
      simplify_rel_tree (Select (r, Sel_and [c1; c2]))
  | Order (r, o) as x ->
      maybe_simplify (fun r -> Order (r, o)) ~child:r x
  | Limit (r, l) as x ->
      maybe_simplify (fun r -> Limit (r, l)) ~child:r x
  | Offset (r, o) as x ->
      maybe_simplify (fun r -> Offset (r, o)) ~child:r x
  | Select (r, c) as x ->
      maybe_simplify (fun r -> Select (r, c)) ~child:r x
  | Table _ as x -> x
  | Rename _ | Project _ | Product _ as x ->
      (* TODO: simplify children *)
      x

and to_sql_aux env = function
    Table (tblname, _, fields) -> (env, select_fields fields ^ " FROM " ^ tblname)

  | Project (t, fields) ->
      let env', text = projection_to_sql_subquery t in
        (merge_environments env env', select_fields fields ^ " FROM " ^ text)

  | Rename (t, fields) ->
      let sel = "SELECT " ^ (String.concat ", "
                     (List.map
                        (fun (dst, src) -> "\"" ^ src ^ "\" AS " ^ "\"" ^ dst ^ "\"")
                        fields)) in
      let env =
        List.fold_left
          (fun s (dst, src) ->
             try
               let target = Idmap.find src env in
                 Idmap.add dst target (Idmap.remove src s)
             with Not_found -> s)
          env fields in
      let env', text = projection_to_sql_subquery t in
        (merge_environments env' env, sel ^ " FROM " ^ text)

  | Product (t, fields1, fields2) ->
      let id1 = new_sql_identifier () in
      let id2 = new_sql_identifier () in
      let prefields1 = List.map (fun f -> (id1, f)) fields1 in
      let prefields2 = List.map (fun f -> (id2, f)) fields2 in
      let sel = select_fields_with_prefix (prefields1 @ prefields2) in
      let env = List.fold_left (fun s f -> Idmap.add f id1 s) env fields1 in
      let env = List.fold_left (fun s f -> Idmap.add f id2 s) env fields2 in
      let env', text = product_to_sql_subquery id1 id2 t in
        (merge_environments env' env, sel ^ " FROM " ^ text)

  | Select (Select (rel, crit2), crit1) ->
      to_sql_aux env (Select (rel, Sel_and [crit1; crit2]))

  | Select (rel, criterion) ->
      let env, text = to_sql_aux env rel in
        (env, text ^ " WHERE " ^ criterion_to_sql rel env criterion)

  | Order (rel, order) ->
      let order_text env l =
        String.concat ", "
          (List.map
             (function
                | Asc (_, field) -> get_full_field_name env field ^ " ASC"
                | Desc(_, field) -> get_full_field_name env field ^ " DESC")
             l) in

      let env, text = to_sql_aux env rel in
        (env, text ^ " ORDER BY " ^ order_text env order)

  | Limit (rel, lim) ->
      let env, text = to_sql_aux env rel in
        (env, text ^ " LIMIT " ^ string_of_int lim)

  | Offset (rel, off) ->
      let env, text = to_sql_aux env rel in
        (env, text ^ " OFFSET " ^ string_of_int off)


and to_sql_subquery env = function
    Table (tblname, _, _) -> (env, tblname)
  | Limit _ | Offset _ | Order _ | Product _ | Project _
  | Rename _ | Select _ as x ->
      let env, text = to_sql_aux env x in
        (env, "(" ^ text ^ ")")

let pack_projection rel f convf =
  { open_projection = fun scope ->
      scope.bind_projection { proj_conv = convf; proj_rel = rel; proj_f = f;
                              proj_to_sql = to_sql_subquery } }


let pack_product fb1 fb2 fc1 fc2 convf relb relc =
  { open_product = fun scope ->
      scope.bind_product
        { prod_conv = convf;
          prod_rel1 = relb; prod_rel2 = relc; prod_f = (fb1, fb2, fc1, fc2);
          prod_to_sql1 = to_sql_subquery; prod_to_sql2 = to_sql_subquery } }

let table name x fields = Table (name, x, fields)

let project f convf fields rel =
  Project (pack_projection (simplify_rel_tree rel) f convf, fields)

let select criterion rel = Select (rel, criterion)

let rename f convf fields rel =
  Rename (pack_projection (simplify_rel_tree rel) f convf, fields)

let product f1a f1b f2a f2b convf fields1 fields2 rel1 rel2 =
  Product ((pack_product f1a f1b f2a f2b convf
              (simplify_rel_tree rel1)
              (simplify_rel_tree rel2)),
           fields1, fields2)

let sel_binop op v ftype field rel =
  Select (rel, Sel_binop (op, v, ftype, field))
let sel_binop' op f1 f2 n1 n2 typef rel =
  Select (rel, Sel_binop' (op, f1, f2, n1, n2, typef))

let order ord_list rel = Order (rel, ord_list)

let limit lim rel = Limit (rel, lim)

let offset off rel = Offset (rel, off)

module ExtStream =
struct
  module Stream =
  struct
    include Stream

    let first s =
      try Stream.next s with Stream.Failure -> raise Not_found

    let to_list s =
      let rec loop l =
        match Stream.peek s with
            Some row -> loop (row :: l)
          | None -> List.rev l
      in loop []

    let to_array s = Array.of_list (to_list s)

    let length s =
      let rec loop s i =
        match Stream.peek s with
            Some _ -> Stream.junk s; loop s (i+1)
          | None -> i
      in loop s 0
  end
end

module Typeconv =
struct
  open Printf

  let encode_string s =
    let len = String.length s in
    let buf = Buffer.create (len * 2) in
      Buffer.add_char buf '\'';
      for i = 0 to len - 1 do
        let c = s.[i] in
          match Char.code c with
            | 39 (* ' *) -> Buffer.add_string buf "''"
            | 92 (* \ *) -> Buffer.add_string buf "\\\\\\\\"
            | b when b >= 0x20 && b <= 0x7e -> Buffer.add_char buf c
            | b -> Buffer.add_string buf (sprintf "\\%03o" b)
      done;
      Buffer.add_char buf '\'';
      Buffer.contents buf

  let oct c = Char.code c - 0x30

  let decode_string s =
    let len = String.length s in
    let buf = Buffer.create len in
    let i = ref 0 in
      while !i < len do
        let c = s.[!i] in
          match c with
            | '\\' ->
                incr i;
                if !i < len then begin
                  let c1 = s.[!i] in
                    match c1 with
                        '\\' -> Buffer.add_char buf '\\'; incr i
                      | '0' .. '3' ->
                          if !i + 2 < len then begin
                            let c2 = s.[!i+1] in
                            let c3 = s.[!i+2] in
                              match c2 with
                                | '0' .. '7' ->
                                    (match c3 with
                                        '0' .. '7' ->
                                          let b = (oct c1 lsl 6) +
                                                  (oct c2 lsl 3) +
                                                  oct c3
                                          in Buffer.add_char buf (Char.chr b);
                                             i := !i + 2
                                      | _ -> invalid_arg "decode_string")
                                | _ -> invalid_arg "decode_string"
                          end
                      | _ -> invalid_arg "decode_string"
                end
            | c -> Buffer.add_char buf c; incr i
      done;
      Buffer.contents buf

  let decode_bool = function
      "true" | "t" -> true
    | "false" | "f" -> false
    | _ -> invalid_arg "decode_bool"

  let encode_bool = function true -> "'t'" | false -> "'f'"
  let encode_int = string_of_int
  let decode_int = int_of_string
  let encode_float = string_of_float
  let decode_float = float_of_string

  let encode_nullable f = function
      Some x -> f x
    | None -> "NULL"

end

module Schema =
struct
  open Printf
  type column_type =
      Boolean | Bytea | Char of int | Date | Double | Int | Real | Serial
      | Smallint | Text | Time of bool | Timestamp of bool | Varchar of int

  type column_attributes = Unique | Indexed

  type column = {
    col_name : string;
    col_type : column_type;
    col_foreign : (string * string) option;
    col_nullable : bool;
    col_prim_key : bool;
    col_attributes : column_attributes list
  }

  type ('a, 'primkey) table = {
    table_relation : 'a; table_name : string; table_columns : column list;
    table_primkey : 'primkey;
  }

  exception Schema_error of string * string

  module Postgresql =
  struct
    let verbose_check = ref false

    let string_of_ctype = function
        Boolean -> "bool" | Bytea -> "bytea"
      | Char n -> sprintf "char(%d)" n
      | Date -> "date" | Double -> "double precision" | Int -> "int"
      | Real -> "real" | Serial -> "serial" | Smallint -> "smallint"
      | Text -> "text"
      | Time b -> if b then "time without timezone" else "time with time zone"
      | Timestamp b ->
          if b then "timestamp without timestampzone"
          else "timestamp with timestamp zone"
    | Varchar n -> sprintf "varchar(%d)" n

    open Postgresql
    let db_col_type = function
        Boolean -> BOOL | Bytea -> BYTEA | Char _ -> CHAR | Date -> DATE
      | Double -> FLOAT8 | Int -> INT4 | Real -> FLOAT4 | Serial -> INT4
      | Smallint -> INT2 | Text -> TEXT
      | Time true -> TIMETZ | Time false -> TIME
      | Timestamp true -> TIMESTAMPTZ | Timestamp false -> TIMESTAMP
      | Varchar _ -> VARCHAR

    let string_of_db_col_type = function
        BOOL -> "bool" | BYTEA -> "bytea" | CHAR -> "char" | DATE -> "date"
      | FLOAT8 -> "double" | INT4 -> "int" | FLOAT4 -> "real"
      | INT2 -> "smallint" | TEXT -> "text" | TIMETZ -> "time with timezone"
      | TIME -> "time" | TIMESTAMPTZ -> "timestamp with timezone"
      | TIMESTAMP -> "timestamp without timezone" | VARCHAR -> "varchar"
      | _ -> invalid_arg "string_of_db_col_type"

    let check schema (db : #Postgresql.connection) =
      let tblname = schema.table_name in
      let doraise x = raise (Schema_error (tblname, x)) in
      let module M = Map.Make(struct type t = string let compare = compare end) in
      let colt_map =
        List.fold_left (fun m c -> M.add c.col_name (db_col_type c.col_type) m)
          M.empty schema.table_columns in
      let r = db#exec ("SELECT * FROM \"" ^ tblname ^ "\" LIMIT 0") in
        if r#status <> Postgresql.Tuples_ok then doraise r#error;
        if r#nfields <> List.length schema.table_columns then
          doraise (sprintf "wrong number of columns. Got %d, expected %d."
                     r#nfields (List.length schema.table_columns));
        for i = 0 to r#nfields - 1 do
          let name = r#fname i in
          let t = r#ftype i in
            try
              let expected = M.find name colt_map in
                if expected <> t then
                  doraise
                    (Printf.sprintf "column type of type %s, expected %s"
                       (string_of_db_col_type t)
                       (string_of_db_col_type expected));
                if !verbose_check then begin
                  Printf.eprintf "column %s.%s of type %s OK\n"
                    schema.table_name name (string_of_db_col_type t)
                end
            with Not_found -> doraise (Printf.sprintf "missing column: %s" name)
        done

    let (<|<) b x = Buffer.add_string b x; b
    let (<<|) b x = Buffer.add_string b x

    let has_attribute attr col = List.mem attr col.col_attributes

    let string_of_column c =
      let b = Buffer.create 64 in
        b <|< "\"" <|< c.col_name <|< "\" " <<| string_of_ctype c.col_type;
        if not c.col_nullable then b <<| " not null";
        if c.col_prim_key then b <<| " primary key";
        if has_attribute Unique c then b <<| " unique";
        Buffer.contents b

    let drop table = sprintf "drop table \"%s\" cascade;" table.table_name

    let generate table =
      let b = Buffer.create 128 in
      let create_index col =
        if has_attribute Indexed col then
          b <<| sprintf "\ncreate %s %s ON %s(%s);\n"
                  (if has_attribute Unique col then "unique index" else "index")
                  (table.table_name ^ "_" ^ col.col_name)
                  table.table_name col.col_name
      in
        (* tables *)
        b <|< "create table \"" <|< table.table_name <|< "\"" <|< "(\n\t" <|<
          String.concat ",\n\t" (List.map string_of_column table.table_columns) <<|
          "\n);";
        (* indexes *)
        List.iter create_index table.table_columns;
        Buffer.contents b

    let constraints table =
      let b = Buffer.create 128 in
      let process_col c = match c.col_foreign with
          None -> ()
        | Some (tbl, col) ->
            b <<| sprintf "ALTER TABLE %s ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s(%s) ON DELETE CASCADE;\n"
                    table.table_name
                    (table.table_name ^ "_fkey_" ^ c.col_name ^ "_" ^ col)
                    c.col_name tbl col
      in List.iter process_col table.table_columns;
         Buffer.contents b
  end
end