open BitVectors;;
open ASM;;
open Util;;
open Parser;;
open Printf;;

exception WrongFormat of string

type intel_hex_entry_type =
    Data
  | End
  | ExtendedSeg
  | ExtendedLinear
;;

type intel_hex_entry =
{
  record_length: byte;
  record_addr: word;
  record_type: intel_hex_entry_type;
  data_field: byte list;
  data_checksum: byte
}
;;

type intel_hex_format = intel_hex_entry list;;

let hex_digit_of_char =
    function
      '0' -> 0 | '1' -> 1 | '2' -> 2
    | '3' -> 3 | '4' -> 4 | '5' -> 5
    | '6' -> 6 | '7' -> 7 | '8' -> 8
    | '9' -> 9 | 'A' -> 10 | 'B' -> 11
    | 'C' -> 12 | 'D' -> 13 | 'E' -> 14
    | 'F' -> 15 | 'a' -> 10 | 'b' -> 11
    | 'c' -> 12 | 'd' -> 13 | 'e' -> 14
    | 'f' -> 15 | _ -> assert false

let intel_hex_entry_type_of_int =
  function
    0 -> Data
  | 1 -> End
  | 2 -> ExtendedSeg
  | 4 -> ExtendedLinear
  | _ -> assert false
;;

let int_of_intel_hex_entry_type =
 function
    Data -> 0
  | End -> 1
  | ExtendedSeg -> 2
  | ExtendedLinear -> 4
;;

let prs_nibble =
         prs_hex_digit >>= 
fun a -> return $ vect_of_int (hex_digit_of_char a) `Four
;;

let prs_byte =
         prs_nibble >>= 
fun a -> prs_nibble >>=
fun b -> return $ mk_byte a b
;;

let prs_word =
         prs_byte >>= 
fun a -> prs_byte >>=
fun b -> return $ mk_word a b
;;

let prs_length = prs_byte;;
let prs_data len = prs_exact len prs_byte
let prs_checksum = prs_byte;;
let prs_addr = prs_word;;

let prs_type =
         prs_hex_digit >>=
fun a -> prs_hex_digit >>=
fun b ->
  let a_as_hex = hex_digit_of_char a in
  let b_as_hex = hex_digit_of_char b in
(*CSC: is next line correct??? *)
  let total = a_as_hex + b_as_hex in
    return $ intel_hex_entry_type_of_int total

let add_bytes v  =
  let r = List.rev v in
  let rec aux (cry, bs) =
    function
      [] -> (cry, bs)
    | hd::tl ->
        aux (half_add hd bs) tl
  in
    aux (false, (vect_of_int 0 `Eight)) r

let calculate_checksum hex_entry =
 let ty = (flip vect_of_int $ `Eight) $ int_of_intel_hex_entry_type hex_entry.record_type in
 let addr1,addr2 = from_word hex_entry.record_addr in
 let _, total = add_bytes (hex_entry.record_length :: addr1 :: addr2 :: ty :: hex_entry.data_field) in
 let _,total = half_add (vect_of_int 1 `Eight) $ complement total in
  total

let checksum_valid hex_entry =
  let total = calculate_checksum hex_entry in
    hex_entry.data_checksum = total

let prs_intel_hex_record =
         prs_char ':'  >>=
fun _ -> prs_length    >>=
fun b -> prs_addr      >>=
fun c -> prs_type      >>=
fun d -> prs_data (int_of_vect b) >>=
fun e -> prs_checksum  >>=
fun f -> prs_eof       >>=
fun _ ->
 let entry =
  { record_length = b;
    record_addr = c;
    record_type = d;
    data_field = e;
    data_checksum = f }
 in
  if checksum_valid entry then
   return entry
  else
   prs_zero
;;

let prs_intel_hex_format =
  prs_sep_by prs_intel_hex_record (prs_char '\n')
;;

let intel_hex_format_of_string s =
  let chars = char_list_of_string s in
    match prs_intel_hex_format chars with
      [] -> None
    | (prs,_)::_ -> Some prs

let string_of_intel_hex_entry entry =
  let b = Buffer.create 655536 in
  let length_string = hex_string_of_vect entry.record_length in
  let addr_string = Printf.sprintf "%04X" (int_of_vect entry.record_addr) in
  let checksum_string = Printf.sprintf "%02X" (int_of_vect entry.data_checksum) in
  let type_string = Printf.sprintf "%02d" (int_of_intel_hex_entry_type entry.record_type) in
  List.iter (Buffer.add_string b)
    [
      ":"; length_string; addr_string; type_string
    ];
  List.iter (fun e -> Buffer.add_string b (hex_string_of_vect e)) entry.data_field;
  Buffer.add_string b checksum_string;
  Buffer.contents b
;;

let string_of_intel_hex_format f =
  let strs = List.map string_of_intel_hex_entry f in
  let rec aux =
    function
      [] -> ""
    | [e] -> e
    | hd::tl -> hd ^ "\n" ^ aux tl
  in
    aux strs

let intel_hex_of_file path =
 let fd = open_in path in
 let rec aux () =
  match try Some (input_line fd) with End_of_file -> None with
     None -> []
   | Some txt ->
      let read = prs_intel_hex_record (Parser.chars_of_string txt) in
      let read =
       match read with
          [x,[]] -> x
        | _ -> raise (WrongFormat txt)
      in
       read::aux ()
 in
  aux ()
;;

let rec load_from mem addr =
 function
    [] -> mem
  | he::tl ->
     load_from (Physical.WordMap.add addr he mem) (snd (BitVectors.half_add addr (BitVectors.vect_of_int 1 `Sixteen))) tl
;;

let process_intel_hex =
 let rec aux mem =
  function
     [] -> assert false
   | he::tl ->
      match he.record_type with
         End -> assert (tl = []); mem
       | Data -> aux (load_from mem he.record_addr he.data_field) tl
       | _ -> assert false
 in
  aux Physical.WordMap.empty
;;

(* DPM: this needs some comment:
     We aim to extract code memory into segmented lists of bytes, with a maximum
     length (chunk_size).  The code memory map has a fixed size (max_addressable)
     on the 8051.  Further, the chunks we extract get segmented when we find an
     unitialized zone in the code memory.
*)
let export_code_memory chunk_size max_addressable code_mem =
  let rec aux chunk address start_address rbuff lbuff =
    if address = max_addressable then
      (start_address, List.rev rbuff)::lbuff
    else if chunk = 0 then
      aux chunk_size address address [] ((start_address, List.rev rbuff)::lbuff)
    else
      let code = Physical.WordMap.find (vect_of_int address `Sixteen) code_mem in
        aux (chunk - 1) (address + 1) start_address (code::rbuff) lbuff
  in
    List.rev (aux chunk_size 0 0 [] [])
;;

let clean_exported_code_memory = List.filter (fun x -> snd x <> [])
;;

let calculate_data_checksum (record_length, record_addr, record_type, data_field) =
  let ty = (flip vect_of_int $ `Eight) $ int_of_intel_hex_entry_type record_type in
  let addr1,addr2 = from_word record_addr in
  let _, total = add_bytes (record_length :: addr1 :: addr2 :: ty :: data_field) in
  let _,total = half_add (vect_of_int 0 `Eight) $ complement total in
    total
;;

let process_exported_code_memory =
  List.map (fun x ->
    let record_length = vect_of_int (List.length (snd x)) `Eight in
    let record_addr = vect_of_int (fst x) `Sixteen in
    let record_type = Data in
    let data_field = snd x in
    let temp_record =
      { record_length = record_length;
        record_addr = record_addr;
        record_type = record_type;
        data_field = data_field;
        data_checksum = zero `Eight
      } in
    { temp_record with data_checksum = calculate_checksum temp_record })
;;

let rec zeros len =
  if len = 0 then
    []
  else
    vect_of_int 0 `Eight :: zeros (len - 1)

let post_process_exported_code_memory intel_hex =
  let reversed = List.rev intel_hex in
  let rec aux hex =
    match hex with
      [] -> []
    | he::tl ->
        if he.record_type = End then
          aux tl
        else if he.record_type = Data then
          if he.data_field = zeros (int_of_vect he.record_length) then
            aux tl
          else
            he::(aux tl)
        else
          tl
  in
    List.rev (aux reversed)

let pack_exported_code_memory chunk_size max_addressable code_mem =
  let export = export_code_memory chunk_size max_addressable code_mem in
  let cleaned = clean_exported_code_memory export in
  let processed = process_exported_code_memory cleaned in
  let postprocessed = post_process_exported_code_memory processed in
  let end_buffer =
    [{ record_length = zero `Eight;
      record_addr = zero `Sixteen;
      record_type = End;
      data_field = [];
      data_checksum = vect_of_int 255 `Eight
    }] in
    postprocessed @ end_buffer
;;

let file_of_intel_hex path fmt =
  let str_fmt = string_of_intel_hex_format fmt in
  let channel = open_out path in
    fprintf channel "%s\n" str_fmt;
    close_out channel
;;