module StringSet = Set.Make (String)

let starts_with ~prefix s =
  let lp = String.length prefix in
  String.length s >= lp && String.sub s 0 lp = prefix

let trim = String.trim

let parse_import_line line =
  let t = trim line in
  if not (starts_with ~prefix:"import" t) then None
  else
    let rest = trim (String.sub t 6 (String.length t - 6)) in
    if String.length rest < 2 || rest.[0] <> '"' then
      Some (Error ("malformed import: " ^ line))
    else
      try
        let q2 = String.index_from rest 1 '"' in
        let path = String.sub rest 1 (q2 - 1) in
        let tail = trim (String.sub rest (q2 + 1) (String.length rest - q2 - 1)) in
        let valid_tail =
          String.equal tail ";"
          || String.equal tail ""
          || (starts_with ~prefix:";" tail && starts_with ~prefix:"//" (trim (String.sub tail 1 (String.length tail - 1))))
        in
        if not valid_tail then Some (Error ("malformed import terminator: " ^ line)) else Some (Ok path)
      with Not_found -> Some (Error ("unterminated import string: " ^ line))

let read_file path =
  let ic = open_in_bin path in
  Fun.protect
    ~finally:(fun () -> close_in ic)
    (fun () -> really_input_string ic (in_channel_length ic))

let resolve_path ~from_dir p = if Filename.is_relative p then Filename.concat from_dir p else p

let load_source_with_imports entry_path =
  let rec load_file stack visited path =
    if List.mem path stack then
      Error ("import cycle detected at: " ^ path)
    else if StringSet.mem path visited then Ok (visited, "")
    else
      let content =
        try Ok (read_file path) with Sys_error msg -> Error ("cannot read " ^ path ^ ": " ^ msg)
      in
      match content with
      | Error _ as e -> e
      | Ok content ->
          let dir = Filename.dirname path in
          let lines = String.split_on_char '\n' content in
          let rec scan imports body = function
            | [] -> Ok (List.rev imports, List.rev body)
            | ln :: tl ->
                (match parse_import_line ln with
                | None -> scan imports (ln :: body) tl
                | Some (Ok p) -> scan (p :: imports) body tl
                | Some (Error e) -> Error e)
          in
          (match scan [] [] lines with
          | Error _ as e -> e
          | Ok (imports, body_lines) ->
              let rec load_imports vis acc = function
                | [] -> Ok (vis, List.rev acc)
                | imp :: tl ->
                    let resolved = resolve_path ~from_dir:dir imp in
                    (match load_file (path :: stack) vis resolved with
                    | Error _ as e -> e
                    | Ok (vis', src) -> load_imports vis' (src :: acc) tl)
              in
              match load_imports visited [] imports with
              | Error _ as e -> e
              | Ok (visited', imported_sources) ->
                  let self_src = String.concat "\n" body_lines in
                  let full = String.concat "\n" (imported_sources @ [ self_src ]) in
                  Ok (StringSet.add path visited', full))
  in
  match load_file [] StringSet.empty entry_path with
  | Error _ as e -> e
  | Ok (_, src) -> Ok src