(* Copyright 2005-2008 Berke DURAK, Inria Rocquencourt, Jaap BOENDER.

This file is part of Dose2.

Dose2 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 3 of the License, or
(at your option) any later version.

Dose2 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, see <http://www.gnu.org/licenses/>. *)

open Io;;
open Napkin;;
module G = Generic_io;;

exception Warning of string;;

module Mmap =
  struct
    type t = Unix.file_descr
    let unmap fd = ()
    let map fd = fd
    let sub fd off lgt =
      ignore (Unix.LargeFile.lseek fd off Unix.SEEK_SET);
      let u = String.create lgt in
      let rec loop i m =
        if m = 0 then
          u
        else
          begin
            let n = Unix.read fd u i m in
            loop (i + n) (m - n)
          end
      in
      loop 0 lgt
  end
;;
let fmt = G.Asc;;
let sf = Printf.sprintf;;
(*** io_field_description, io_field *)
let io_field_description = io_pair io_int (io_list (io_triple io_string io_int64 io_int));;
let io_field = io_pair io_string (io_list io_string);;
let io_log = io_quadruple io_float io_int io_string io_string;;
(* ***)
let field_name_prefix = "field.";;
(*** sanitize_field_name, unsanitize_field_name *)
let sanitize_field_name b fn = Util.sanitize_filename ~buffer:b ~prefix:field_name_prefix fn;;
let unsanitize_field_name b fn = Util.unsanitize_filename ~buffer:b ~prefix:field_name_prefix fn;;
(* ***)
(*** packages_cat,lifetime_cat,count_cat *)
let packages_cat = "pkg";;
let lifetime_cat = "lf";;
let population_cat = "pop";;
let type_cat = "type";;
(* ***)
let separator = '\t';;
(*** make_key *)
let make_key b category elements =
  Buffer.clear b;
  List.iter
    begin fun k ->
      Buffer.add_char b separator;
      Buffer.add_string b k
    end
    (category :: elements);
  Buffer.contents b
;;
(* ***)
(*** make_lifetime_key *)
let make_lifetime_key ~buffer ~archive ~unit_name ~version ~architecture =
  make_key buffer lifetime_cat [archive;unit_name;version;architecture]
;;
(* ***)
let make_package_key ~buffer ~unit_name ~version ~architecture =
  make_key buffer packages_cat [unit_name;version;architecture]
;;
(* ***)
(*** decompose_key *)
let decompose_key k =
  match Util.split_at separator k with
  | category :: elements -> (category, elements)
  | _ -> invalid_arg "decompose_key"
;;
(* ***)
(*** get_field *)
let rec get_field key = function
  | [] -> raise Not_found
  | (key',v1,v2)::rest ->
      if Util.lowercase_compare key key' then
        (v1,v2)
      else
        get_field key rest
;;
(* ***)
(*** get_field_first *)
let rec get_field_first key = function
  | [] -> raise Not_found
  | (key',v1,_)::rest ->
      if Util.lowercase_compare key key' then
        v1
      else
        get_field_first key rest
;;
(* ***)
(*** Out *)
module Out =
  struct
    type t = {
      db_base : string;
      db_timestamp : float;
      mutable db_count : int;
      mutable db_population : int;
			mutable db_type : string;
      db_dbm : Dbm.t;
      db_fields : (string, out_channel) Hashtbl.t;
      db_buffer : Buffer.t;
      db_log : out_channel
    }

    (*** count_population *)
    let count_population dbm =
      let count = ref 0 in
      Dbm.iter
        begin fun key data ->
          match decompose_key key with
          | (category, _) when category = packages_cat -> incr count
          | _ -> ()
        end
        dbm;
      !count
    ;;
    (* ***)
    (*** open_out *)
    let open_out ?set_type base =
      if not (Sys.file_exists base) then
        Unix.mkdir base 0o755;
      let log_file =
        Filename.concat base (Printf.sprintf "log.%s" (Util.string_of_iso_8601 (Util.today ())))
      in
      let dbm = Dbm.opendbm (Filename.concat base "index") [Dbm.Dbm_rdwr; Dbm.Dbm_create] 0o644 in
      let population =
        try
          int_of_string (Dbm.find dbm population_cat)
        with
        | Not_found -> count_population dbm
			in
			let db_type =
				try
					Dbm.find dbm type_cat
				with
				| Not_found -> "unknown"
      in
			let type_to_set = match set_type with
				None -> db_type
			| Some s -> begin
					if db_type = s then s
					else if db_type = "unknown" then s
					else raise (Failure (Printf.sprintf ("Type in database and type to set do not match (%s vs %s)") db_type s))
				end
			in
      { db_base = base;
        db_timestamp = Unix.gettimeofday ();
        db_log = open_out_gen [Open_wronly;Open_append;Open_creat;Open_binary] 0o644 log_file;
        db_count = 0;
        db_population = population;
				db_type = type_to_set;
        db_dbm = dbm;
        db_buffer = Buffer.create 256;
        db_fields = Hashtbl.create 256 }
    ;;
    (* ***)
    (*** get_field_file *)
    let get_field_file db field =
      try
        Hashtbl.find db.db_fields field
      with
      | Not_found ->
         let field' = String.lowercase field in
         try
           Hashtbl.find db.db_fields field'
         with
         | Not_found ->
           (*let dbm = Dbm.opendbm ("field."^field^".database") [Dbm.Dbm_rdwr; Dbm.Dbm_create] 0o644 in*)
           let fn = sanitize_field_name db.db_buffer field' in
           let oc = open_out_gen [Open_wronly;Open_append;Open_creat;Open_binary] 0o644 (Filename.concat db.db_base fn) in
           LargeFile.seek_out oc (LargeFile.out_channel_length oc);
           Hashtbl.add db.db_fields field oc;
           Hashtbl.add db.db_fields field' oc;
           oc
    ;;
    (* ***)
    (*** dbm_replace *)
    let dbm_replace db key data =
      (*Printf.eprintf "%f %d %S (%d)\n%!" db.db_timestamp db.db_count key (String.length data);*)
      G.save_to_channel ~header:false ~fmt io_log (db.db_timestamp,db.db_count,key,data) db.db_log;
      output_char db.db_log '\n';
      db.db_count <- 1 + db.db_count;
      Dbm.replace db.db_dbm key data
    ;;
    (* ***)
    (*** close_out *)
    let close_out db =
      dbm_replace db type_cat db.db_type;
      close_out db.db_log;
      Dbm.close db.db_dbm;
      Hashtbl.iter (fun _ oc -> close_out oc) db.db_fields
    ;;
    (* ***)
    (*** rebuild_index *)
    (* Rebuild the index from log files. *)
    let rebuild_index base =
      let fn = Filename.concat base "index" in
      let attempt f = try ignore (Lazy.force f) with _ -> () in
      attempt (lazy (Unix.unlink (fn^".dir")));
      attempt (lazy (Unix.unlink (fn^".pag")));
      let dbm = Dbm.opendbm (Filename.concat base "index") [Dbm.Dbm_rdwr; Dbm.Dbm_create] 0o644 in
      let dir = Unix.opendir base in
      try
        let times = Hashtbl.create 16384 in
        while true do
          let fn = Unix.readdir dir in
          if Util.is_prefix "log." fn then
            begin
              Printf.printf "Processing %S...\n%!" fn;
              let ic = open_in_bin (Filename.concat base fn) in
              let n = in_channel_length ic in
              begin
                try
                  let i = G.in_from_channel ~header:false ~fmt ic in
                  while pos_in ic < n do
                    let (timestamp,count,key,data) = read io_log i in
                    Printf.printf "%f %d %S (%d)\n%!" timestamp count key (String.length data);
                    try
                      let old_timestamp = Hashtbl.find times key in
                      if (timestamp,count) > old_timestamp then
                        begin
                          Dbm.replace dbm key data;
                          Hashtbl.replace times key (timestamp,count)
                        end
                      else
                        ()
                    with
                    | Not_found ->
                        Dbm.replace dbm key data;
                        Hashtbl.add times key (timestamp,count)
                  done;
                  Printf.printf "Processed %S OK.\n%!" fn
                with
                | x -> Printf.printf "Error: %s\n%!" (Printexc.to_string x)
              end;
              close_in ic
            end
        done
      with
      | End_of_file ->
          Unix.closedir dir;
          Dbm.close dbm
    ;;
    (* ***)
    (*** add_lifetime *)
    let add_lifetime db ~archive ~unit_name ~version ~architecture ~day =
      let key = make_lifetime_key ~buffer:db.db_buffer ~archive ~unit_name ~version ~architecture in
      let lf =
        try
          G.from_string ~header:false ~fmt Lifetime.io_lifetime (Dbm.find db.db_dbm key)
        with
        | Not_found -> Lifetime.empty
      in
      let lf = Lifetime.add_day lf day in
      dbm_replace db key (G.to_string ~header:false ~fmt Lifetime.io_lifetime lf)
    ;;
    (* ***)
    (*** add_package *)
    let add_package db ~archive ~day fields =
      let unit_name         = get_field_first "package" fields
      and version      = get_field_first "version" fields
      and architecture = get_field_first "architecture" fields
      in
      add_lifetime db ~archive ~unit_name ~version ~architecture ~day;
      let key = make_package_key ~buffer:db.db_buffer ~unit_name ~version ~architecture in
      try
        ignore (Dbm.find db.db_dbm key)
      with
      | Not_found ->
        let package_id = db.db_population in
        db.db_population <- 1 + db.db_population;
        let mp =
          List.map
            begin fun (field, first, rest) ->
              let oc = get_field_file db field in
              let x = (first, rest) in
              let u = G.to_string ~header:false ~fmt io_field x in
              let m = String.length u in
              let entry = (field, LargeFile.pos_out oc, m) in
              output_string oc u;
              output_char oc '\n';
              entry
            end
            fields
        in
        let mp' = G.to_string ~header:false ~fmt io_field_description (package_id, mp) in
        dbm_replace db key mp';
        dbm_replace db population_cat (sf "%d" db.db_population)
    ;;
    (* ***)
  end
;;
(* ***)
(*** In *)
module In =
  struct
    type t = {
      db_base : string;
      db_dbm : Dbm.t;
      db_fields : (string, Unix.file_descr * Mmap.t) Hashtbl.t;
      db_known_fields : string list;
      db_buffer : Buffer.t
    };;

    type package = int * (string * int64 * int) list;;

    (*** known_fields *)
    let known_fields db = db.db_known_fields;;
    (* ***)
    (*** scan_fields *)
    let scan_fields base buffer =
      let dir = Unix.opendir base in
      let rec loop result =
        match
          try
            Some(Unix.readdir dir)
          with
          | End_of_file ->
              Unix.closedir dir;
              None
        with
        | None -> result
        | Some name ->
            if Util.is_prefix field_name_prefix name then
              begin
                let field = unsanitize_field_name buffer name in
                loop (field :: result)
              end
            else
              loop result
      in
      loop []
    ;;
    (* ***)
    (*** open_in *)
    let open_in base =
      let buffer = Buffer.create 256 in
      { db_base = base;
        db_dbm = Dbm.opendbm (Filename.concat base "index") [Dbm.Dbm_rdonly] 0o644;
        db_buffer = buffer;
        db_known_fields = scan_fields base buffer;
        db_fields = Hashtbl.create 256 }
    ;;
    (* ***)
    (*** close_in *)
    let close_in db =
      Dbm.close db.db_dbm;
      List.iter
        begin fun field ->
          try
            let (fd,mp) = Hashtbl.find db.db_fields field in
            Mmap.unmap mp;
            Unix.close fd
          with
          | Not_found -> ()
        end
        db.db_known_fields
    ;;
    (* ***)
		(*** count_population *)
		let count_population db =
			int_of_string (Dbm.find db.db_dbm population_cat)
		;;
		(* ***)
		let get_type db =
			try
				Dbm.find db.db_dbm type_cat
			with Not_found -> "unknown"
		;;
    (*** iter *)
    let iter
      ?(package = (fun ~unit_name ~version ~architecture _ -> ()))
      ?(lifetime = (fun ~archive ~unit_name ~version ~architecture _ -> ()))
      db
      =
      Dbm.iter
        begin fun key data ->
          match decompose_key key with
          | (category, [unit_name;version;architecture]) when category = packages_cat ->
              let p = G.from_string ~header:false ~fmt io_field_description data in
              package ~unit_name ~version ~architecture p
          | (category, [archive;unit_name;version;architecture]) when category = lifetime_cat ->
              let lf = G.from_string ~header:false ~fmt Lifetime.io_lifetime data in
              lifetime ~archive ~unit_name ~version ~architecture lf
          | _ -> ()
        end
        db.db_dbm
    (* ***)
    (*** find_package *)
    let find_package db ~unit_name ~version ~architecture =
      let key = make_package_key ~buffer:db.db_buffer ~unit_name ~version ~architecture in
      let data = Dbm.find db.db_dbm key in
      G.from_string ~header:false ~fmt io_field_description data
    ;;
    (* ***)
    (*** find_lifetime *)
    let find_lifetime db ~architecture ~archive ~unit_name ~version =
      let key = make_lifetime_key ~buffer:db.db_buffer ~archive ~unit_name ~version ~architecture in
      let data = Dbm.find db.db_dbm key in
      G.from_string ~header:false ~fmt Lifetime.io_lifetime data
    ;;
    (* ***)
    (*** get_field_map *)
    let get_field_map db field =
      try
        let (_,mp) = Hashtbl.find db.db_fields field in mp
      with
      | Not_found ->
         let field' = String.lowercase field in
         try
           let (_,mp) = Hashtbl.find db.db_fields field' in mp
         with
         | Not_found ->
           let fn = sanitize_field_name db.db_buffer field' in
           let fd = Unix.openfile (Filename.concat db.db_base fn) [Unix.O_RDONLY] 0o644 in
           let mp = Mmap.map fd in
           Hashtbl.add db.db_fields field (fd,mp);
           Hashtbl.add db.db_fields field' (fd,mp);
           mp
    ;;
    (* ***)
    (*** extract_field *)
    let extract_field db field offset length =
      let mp = get_field_map db field in
      let x = Mmap.sub mp offset length in
      try
        G.from_string ~header:false ~fmt io_field x
      with
      | x ->
          (*Printf.printf "SHIT: %s\n%!" (Printexc.to_string x);
          Printf.printf "field=%s offset=%Ld length=%d\n%!" field offset length;*)
          raise x
    ;;
    (* ***)
    (*** get_package_id *)
    let get_package_id db (package_id,_) = package_id;;
    (* ***)
    (*** get_field *)
    let get_field db (_,fds) field =
      let (offset,length) = get_field field fds in
      extract_field db field offset length
    ;;
    (* ***)
    (*** iter_over_fields *)
    let iter_over_fields f db (_,pkg) =
      List.iter
        begin fun (field,offset,length) ->
          let (first,rest) = extract_field db field offset length in
          f (field,first,rest)
        end
        pkg
    ;;
    (* ***)
    (*** get_field_first *)
    let get_field_first db fds field =
      let (x,_) = get_field db fds field in x
    ;;
    (* ***)
  end
;;
(* ***)

let napkin_of_fields =
  let rex0 = Pcre.regexp ~study:true "\\s*(\\S+)\\s*"
  and rex1 = Pcre.regexp ~study:true "\\s*,\\s*"
  and rex2 = Pcre.regexp ~study:true "\\s*\\|\\s*"
  and rex3 = Pcre.regexp ~study:true "(\\S+)\\s*\\(\\s*(<|<<|<=|>=|>>|>|=)\\s*([0-9a-zA-Z.:+-~]+)\\s*\\)"
  and rex4 = Pcre.regexp ~study:true "^\\s*(\\S+)\\s*\\((\\S+)\\)$"
  in
  fun assoc ->
    let chump x =
      let ss = Pcre.exec ~rex:rex0 x in
      Pcre.get_substring ss 1
    in
    let (pk_unit, pk_version) =
      let u = chump (assoc "Package") in
      try
        let v = chump (assoc "Version") in
        (u,v)
      with
      | Not_found -> raise (Warning(sf "Package %S ignored because it has no version." u))
    in
    let get k =
      try
        Pcre.split ~rex:rex1 (assoc k)
      with
      | Not_found -> []
    in
    (*** Source *)
    let pk_source =
      try
        let ss = assoc "Source" in
        try
          let suv = Pcre.exec ~rex:rex4 ss in
          (Pcre.get_substring suv 1, Pcre.get_substring suv 2)
        with
        | Not_found -> (chump ss, pk_version)
      with
			| Not_found -> (pk_unit, pk_version)
    in
    (* ***)
    (*** Get an integer value *)
    let get_int64 field =
      try
        Int64.of_string (assoc field)
      with
      | Not_found|Invalid_argument _ -> 0L
    in
    (* ***)
    (*** Get a boolean value *)
    let get_bool field =
      try
        assoc field = "yes"
      with
      | Not_found -> false
    in
    (* ***)
    let pk_architecture = assoc "Architecture" in
    (*** Collect a conjunctive field *)
    let collect_conjunctive field =
      let stuff = get field in
      List.map (fun x ->
          try
            let ss = Pcre.exec ~rex:rex3 x in
            let (pn,op,v) = (Pcre.get_substring ss 1, Pcre.get_substring ss 2, Pcre.get_substring ss 3) in
            match op with
            | "<"     -> Unit_version (pn, Sel_LT v) (* XXX *)
            | "<<"    -> Unit_version (pn, Sel_LT v)
            | "<="    -> Unit_version (pn, Sel_LEQ v)
            | "="     -> Unit_version (pn, Sel_EQ v)
            | ">="    -> Unit_version (pn, Sel_GEQ v)
            | ">"     -> Unit_version (pn, Sel_GT v)
            | ">>"|_  -> Unit_version (pn, Sel_GT v)
          with
          | Not_found -> Unit_version (chump x,  Sel_ANY)) stuff
    in
    (* ***)
    (*** Collect disjunctive field *)
    let collect_disjunctive field =
      let stuff = get field in
      List.map (fun d ->
        let l = Pcre.split ~rex:rex2 d in
        List.map
          (fun x ->
            try
              let ss = Pcre.exec ~rex:rex3 x in
              let (pn,op,v) = (Pcre.get_substring ss 1, Pcre.get_substring ss 2, Pcre.get_substring ss 3) in
              match op with
              | "<"     -> Unit_version (pn, Sel_LT v)
              | "<<"    -> Unit_version (pn, Sel_LT v)
              | "<="    -> Unit_version (pn, Sel_LEQ v)
              | "="     -> Unit_version (pn, Sel_EQ v)
              | ">="    -> Unit_version (pn, Sel_GEQ v)
              | ">"     -> Unit_version (pn, Sel_GT v)
              | ">>"|_  -> Unit_version (pn, Sel_GT v)
            with
            | Not_found ->
                Unit_version (chump x, Sel_ANY)) l) stuff
    in
    (* ***)
    (*** Essential, build-essential *)
    { pk_extra           = ();
      pk_unit            = pk_unit;
      pk_version         = pk_version;
      pk_architecture    = pk_architecture;
      pk_source          = pk_source;
      pk_essential       = get_bool "Essential";
      pk_build_essential = get_bool "Build-Essential";
      pk_provides        = collect_conjunctive "Provides";
      pk_size            = get_int64 "Size";
      pk_installed_size  = get_int64 "Installed-Size";
      pk_conflicts       = collect_conjunctive "Conflicts";
      pk_breaks          = collect_conjunctive "Breaks";
      pk_replaces        = collect_conjunctive "Replaces";
      pk_depends         = collect_disjunctive "Depends";
      pk_pre_depends     = collect_disjunctive "Pre-Depends";
      pk_suggests        = collect_disjunctive "Suggests";
      pk_recommends      = collect_disjunctive "Recommends";
      pk_enhances        = collect_disjunctive "Enhances";
    }
;;
