(* Defen(estrator) - A Freesite Insertion and Management Tool       *)
(* by Travis Bemann and Eric Norige                                 *)
(*                                                                  *)
(* 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; either     *)
(* version 2 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.              *)

exception Error of string
exception Type_error

type t =
    Group of t list
  | Name of string
  | String of string
  | Int of int
  | Float of float

type token =
    Token_open_paren
  | Token_close_paren
  | Token_name of string
  | Token_string of string
  | Token_int of int
  | Token_float of float

let tokenize_name data index =
  let len = String.length data in
  let rec step index =
    if index < len then
      match String.get data index with
	'A' .. 'Z' | 'a' .. 'z' | '0' .. '9' | '!' | '$' | '%' | '&' |
	'*' | '+' | '-' | '.' | '/' | ':' | '<' | '=' | '>' | '?' |
	'@' | '^' |'_' | '~' -> step (index + 1)
      | ' ' | '\t' | '\n' | '\r' | '(' | ')' -> index
      | _ -> raise (Error "Unexpected character")
    else index
  in
  let end_index = step index in
  let name_len = end_index - index in
  let buf = String.create name_len in
  String.blit data index buf 0 name_len;
  (Token_name buf, end_index)

let tokenize_string data index =
  let len = String.length data in
  let buf = Buffer.create 0 in
  let rec step index =
    if index < len then
      match String.get data index with
        '\\' -> step_escape (index + 1)
      | '"' -> (Token_string (Buffer.contents buf), index + 1)
      | ch ->
	  Buffer.add_char buf ch;
	  step (index + 1)
    else raise (Error "Unclosed string - a")
  and step_escape index =
    if index < len then
      match String.get data index with
	't' ->
	  Buffer.add_char buf '\t';
	  step (index + 1)
      | 'n' ->
	  Buffer.add_char buf '\n';
	  step (index + 1)
      | 'r' ->
	  Buffer.add_char buf '\r';
	  step (index + 1)
      | 'b' ->
	  Buffer.add_char buf '\b';
	  step (index + 1)
      | 'x' ->
	  step_hex_escape (index + 1)
      | '\\' ->
	  Buffer.add_char buf '\\';
	  step (index + 1)
      | _ -> raise (Error "Unexpected escape character")
    else raise (Error "Unclosed string - b")
  and step_hex_escape index =
    if index < len - 1 then begin
      Buffer.add_char buf
	(char_of_int (Strutil.int_of_hex (String.sub data index 2)));
      step (index + 2)
    end else raise (Error "Unclosed string - c")
  in step index

let tokenize_number data index =
  let len = String.length data in
  let rec step index is_float =
    if index < len then
      match String.get data index with
	'0' .. '9' -> step (index + 1) is_float
      | '.' ->
	  if not is_float then
	    step (index + 1) true
	  else raise (Error "Invalid number")
      | ' ' | '\t' | '\n' | '\r' | '(' | ')' -> (index, is_float)
      | _ -> raise (Error "Unexpected character")
    else (index, is_float)
  in
  let end_index, is_float = step index false in
  let data = String.sub data index (end_index - index) in
  if not is_float then
    (Token_int (int_of_string data), end_index)
  else (Token_float (float_of_string data), end_index)

let tokenize data =
  let len = String.length data in
  let rec step index tokens_rev =
    if index < len then
      let ch = String.get data index in
      match ch with
	' ' | '\t' | '\n' | '\r' -> step (index + 1) tokens_rev
      | _ ->
	  let token, index =
	    match ch with
	      '(' -> (Token_open_paren, index + 1)
	    | ')' -> (Token_close_paren, index + 1)
	    | 'A' .. 'Z' | 'a' .. 'z' | '!' | '$' | '%' | '&' | '*' |
	      '+' | '-' | '.' | '/' | ':' | '<' | '=' | '>' | '?' |
	      '@' | '^' | '_' | '~' -> tokenize_name data index
	    | '"' -> tokenize_string data (index + 1)
	    | '0' .. '9' -> tokenize_number data index
	    | _ -> raise (Error "Bad file format")
	  in step index (token :: tokens_rev)
    else List.rev tokens_rev
  in step 0 []

let rec parse_sexpr_group tokens items_rev =
  match tokens with
    Token_close_paren :: rest -> (Group (List.rev items_rev), rest)
  | _ ->
      let sexpr, rest = parse_sexpr tokens in
      parse_sexpr_group rest (sexpr :: items_rev)

and parse_sexpr = function
    Token_open_paren :: rest -> parse_sexpr_group rest []
  | Token_name name :: rest -> (Name name, rest)
  | Token_string string :: rest -> (String string, rest)
  | Token_int int :: rest -> (Int int, rest)
  | Token_float float :: rest -> (Float float, rest)
  | Token_close_paren :: _ -> raise (Error "Unexpected token")
  | [] -> raise (Error "Unexpected end of data")

let sexprs_of_tokens tokens =
  let rec step tokens sexprs_rev =
    match tokens with
      [] -> List.rev sexprs_rev
    | _ ->
	let sexpr, rest = parse_sexpr tokens in
	match sexpr with
	  Group _ -> step rest (sexpr :: sexprs_rev)
	| _ -> raise (Error "Only groups may exist at the top level")
  in step tokens []

let of_string data = sexprs_of_tokens (tokenize data)

let bound_column = 71
let group_indent = 3

let add_indent buf indent =
  Buffer.add_char buf '\n';
  for i = 0 to indent - 1 do
    Buffer.add_char buf ' '
  done

let prettyprint_data data buf indent column first =
  let len = String.length data in
  if not first then
    if (column + len) + 1 < bound_column then begin
      Buffer.add_char buf ' ';
      Buffer.add_string buf data;
      (column + len) + 1
    end else begin
      add_indent buf indent;
      Buffer.add_string buf data;
      indent + len
    end
  else
    if column + len < bound_column then begin
      Buffer.add_string buf data;
      column + len
    end else begin
      add_indent buf indent;
      Buffer.add_string buf data;
      indent + len
    end

let prettyprint_string string buf indent column first =
  let len = String.length string in
  let sbuf = Buffer.create (len + 1) in
  let rec step index =
    if index < len then begin
      let ch = String.get string index in
      begin match ch with
	'\t' -> Buffer.add_string sbuf "\\t"
      | '\n' -> Buffer.add_string sbuf "\\n"
      | '\r' -> Buffer.add_string sbuf "\\r"
      | '\b' -> Buffer.add_string sbuf "\\b"
      | '\\' -> Buffer.add_string sbuf "\\\\"
      | _ ->
	  let code = int_of_char ch in
	  if (code >= 0x20) && (code <> 0x7F) then
	    Buffer.add_char sbuf ch
	  else Printf.bprintf sbuf "\\x%02X" code
      end;
      step (index + 1)
    end else begin
      Buffer.add_char sbuf '"';
      Buffer.contents sbuf
    end
  in
  Buffer.add_char sbuf '"';
  prettyprint_data (step 0) buf indent column first

let prettyprint_int int buf indent column first =
  prettyprint_data (string_of_int int) buf indent column first

let prettyprint_float float buf indent column first =
  prettyprint_data (string_of_float float) buf indent column first

let rec prettyprint sexpr buf indent column first =
  match sexpr with
    Group items -> prettyprint_group items buf indent column first
  | Name name -> prettyprint_data name buf indent column first
  | String string -> prettyprint_string string buf indent column first
  | Int int -> prettyprint_int int buf indent column first
  | Float float -> prettyprint_float float buf indent column first

and prettyprint_group items buf indent start_column first =
  let column = prettyprint_data "(" buf indent start_column first in
  let rec step items column first =
    match items with
      item :: rest ->
	let column = prettyprint item buf
	    (start_column + group_indent) column first
	in step rest column false
    | [] ->
	Buffer.add_char buf ')';
	column + 1
  in step items column true

let to_string sexprs =
  let buf = Buffer.create 0 in
  List.iter
    begin fun sexpr ->
      let _ = prettyprint sexpr buf 0 0 true in
      Buffer.add_string buf "\n\n"
    end
    sexprs;
  Buffer.contents buf

let get_list = function
    Group items -> items
  | _ -> raise Type_error

let get_string = function
    String string -> string 
  | _ -> raise Type_error

let get_int = function
    Int int -> int
  | _ -> raise Type_error

let get_float = function
    Float float -> float
  | _ -> raise Type_error

let expand_name = function
    Name name -> name
  | _ -> raise Type_error

let rec get_toplevel_name sexprs name =
  match sexprs with
    Group (Name item :: data :: _) :: rest ->
      if item = name then
	data
      else get_toplevel_name rest name
  | _ :: rest -> get_toplevel_name rest name
  | [] -> raise Not_found

let get_name sexpr name =
  match sexpr with
    Group sexprs -> get_toplevel_name sexprs name
  | _ -> raise Not_found

let rec get_toplevel_name_mult sexprs name =
  match sexprs with
    Group (Name item :: data) :: rest ->
      if item = name then
	data
      else get_toplevel_name_mult rest name
  | _ :: rest -> get_toplevel_name_mult rest name
  | [] -> raise Not_found

let get_name_mult sexpr name =
  match sexpr with
    Group sexprs -> get_toplevel_name_mult sexprs name
  | _ -> raise Not_found

let rec prepend_rev x y =
  match x with
    item :: rest -> prepend_rev rest (item :: y)
  | [] -> y

let set_toplevel_name sexprs name data =
  let rec step sexprs past =
    match sexprs with
      Group (Name item :: _) as cur :: rest ->
	if item = name then
	  prepend_rev rest (Group [Name name; data] :: past)
	else step rest (cur :: past)
    | cur :: rest -> step rest (cur :: past)
    | [] -> Group [Name name; data] :: past
  in step sexprs []

let set_name sexpr name data =
  match sexpr with
    Group sexprs -> set_toplevel_name sexprs name data
  | _ -> raise Type_error

let set_toplevel_name_mult sexprs name data =
  let rec step sexprs past =
    match sexprs with
      Group (Name item :: _) as cur :: rest ->
	if item = name then
	  prepend_rev rest (Group (Name name :: data) :: past)
	else step rest (cur :: past)
    | cur :: rest -> step rest (cur :: past)
    | [] -> Group (Name name :: data) :: past
  in step sexprs []

let set_name_mult sexpr name data =
  match sexpr with
    Group sexprs -> set_toplevel_name_mult sexprs name data
  | _ -> raise Type_error