///|
suberror CliError {
  Message(String)
} derive(Eq, @debug.Debug)

///|
impl Show for CliError with fn output(self, logger) {
  match self {
    Message(message) => logger.write_string(message)
  }
}

///|
fn[T] fail(message : String) -> T raise CliError {
  raise CliError::Message(message)
}

///|
enum QueryMode {
  Or
  And
  Phrase
  QueryString
} derive(Eq, @debug.Debug)

///|
enum OutputFormat {
  Text
  Json
} derive(Eq, @debug.Debug)

///|
struct IndexReport {
  document_count : Int
  generation : Int
  segment_count : Int
}

///|
struct StoredFieldReport {
  name : String
  values : ReadOnlyArray[String]
}

///|
struct HighlightValueReport {
  value_index : Int
  fragment : @lib.HighlightFragment
}

///|
struct HighlightFieldReport {
  name : String
  values : ReadOnlyArray[HighlightValueReport]
}

///|
struct SearchHitReport {
  segment_ord : Int
  doc_id : Int
  score : Double
  fields : ReadOnlyArray[StoredFieldReport]
  highlights : ReadOnlyArray[HighlightFieldReport]
}

///|
struct SearchReport {
  field : String
  query : String
  hits : ReadOnlyArray[SearchHitReport]
}

///|
struct FieldReport {
  id : Int
  name : String
  indexed : Bool
  stored : Bool
  tokenizer : String
}

///|
struct InspectReport {
  generation : Int
  segment_count : Int
  document_count : Int
  fields : ReadOnlyArray[FieldReport]
}

///|
fn validate_fields(field_names : Array[String]) -> Unit raise CliError {
  guard field_names.length() > 0 else {
    fail("at least one --field is required")
  }
  for index in 0.. 0 else {
      fail("field names must not be empty")
    }
    for earlier in 0.. @lib.Schema raise CliError {
  validate_fields(field_names)
  let tokenizers = @lib.TokenizerManager::with_defaults()
  guard tokenizers.contains(tokenizer_name) else {
    fail("unknown tokenizer: \{tokenizer_name}")
  }
  let builder = @lib.SchemaBuilder::new()
  for name in field_names {
    ignore(
      builder.add_text_field(
        name,
        @lib.TextOptions::new(true, true).with_tokenizer(tokenizer_name),
      ),
    )
  }
  builder.build()
}

///|
fn add_json_field_values(
  document : @lib.Document,
  field_id : @lib.FieldId,
  field_name : String,
  value : Json,
  line_number : Int,
) -> Unit raise CliError {
  match value {
    String(text) => document.add_text(field_id, text)
    Array(values) =>
      for value in values {
        match value {
          String(text) => document.add_text(field_id, text)
          _ =>
            fail(
              "JSONL line \{line_number} field '\{field_name}' must be a string or an array of strings",
            )
        }
      }
    _ =>
      fail(
        "JSONL line \{line_number} field '\{field_name}' must be a string or an array of strings",
      )
  }
}

///|
fn parse_jsonl(
  text : String,
  field_names : Array[String],
) -> Array[@lib.Document] raise CliError {
  let documents : Array[@lib.Document] = []
  let mut line_number = 0
  for raw_line in text.split("\n") {
    line_number += 1
    let line = raw_line.to_owned().trim().to_owned()
    if line.length() == 0 {
      continue
    }
    let json = @json.parse(line) catch {
      error => fail("invalid JSON on line \{line_number}: \{error}")
    }
    guard json is Object(object) else {
      fail("JSONL line \{line_number} must be an object")
    }
    let document = @lib.Document::new()
    for field_index in 0..
          add_json_field_values(
            document,
            @lib.FieldId::new(field_index),
            field_name,
            value,
            line_number,
          )
        None => ()
      }
    }
    documents.push(document)
  }
  guard documents.length() > 0 else {
    fail("JSONL input contains no documents")
  }
  documents
}

///|
fn[D : @lib.DirectoryV2] index_jsonl(
  directory : D,
  text : String,
  field_names : Array[String],
  tokenizer_name : String,
) -> IndexReport raise CliError {
  let schema = build_schema(field_names, tokenizer_name)
  let documents = parse_jsonl(text, field_names)
  let tokenizers = @lib.TokenizerManager::with_defaults()
  let segment_writer = @lib.SegmentWriter::with_schema_and_tokenizers(
    schema, tokenizers,
  ) catch {
    error => fail(error.to_string())
  }
  for document in documents {
    ignore(segment_writer.add_document(document))
  }
  let writer = @lib.IndexWriter::new(directory)
  writer.commit(segment_writer.finish()) catch {
    error => fail(error.to_string())
  }
  writer.close()
  let reader = @lib.IndexReader::open(directory) catch {
    error => fail(error.to_string())
  }
  let report = {
    document_count: documents.length(),
    generation: reader.generation() catch {
      error => fail(error.to_string())
    },
    segment_count: reader.segment_count() catch {
      error => fail(error.to_string())
    },
  }
  reader.close()
  report
}

///|
fn stored_fields(
  schema : @lib.Schema,
  document : @lib.StoredDocument,
) -> ReadOnlyArray[StoredFieldReport] {
  let fields : Array[StoredFieldReport] = []
  for field_index in 0.. 0 {
      let name = match schema.field_name(field_id) {
        Some(name) => name
        None => abort("schema field disappeared while rendering a hit")
      }
      fields.push({ name, values })
    }
  }
  ReadOnlyArray::from_array(fields)
}

///|
fn[D : @lib.Directory] search_index(
  directory : D,
  field_name : String,
  query_text : String,
  limit : Int,
  mode : QueryMode,
) -> SearchReport raise CliError {
  search_index_internal(directory, field_name, query_text, limit, mode, None)
}

///|
fn highlight_field(
  highlighter : @lib.Highlighter,
  field_id : @lib.FieldId,
  field_name : String,
  query_text : String,
  mode : QueryMode,
  fragment_size : Int,
  document : @lib.StoredDocument,
) -> ReadOnlyArray[HighlightFieldReport] raise CliError {
  let options = @lib.HighlightOptions::new()
    .with_max_chars(fragment_size)
    .with_mode(
      match mode {
        Phrase => @lib.HighlightMode::Phrase
        Or | And | QueryString => @lib.HighlightMode::Terms
      },
    )
  let values : Array[HighlightValueReport] = []
  let stored_values = document.texts_for(field_id)
  for value_index in 0.. fail(error.to_string())
    }
    match fragment {
      Some(fragment) => values.push({ value_index, fragment })
      None => ()
    }
  }
  if values.length() == 0 {
    ReadOnlyArray::from_array([])
  } else {
    ReadOnlyArray::from_array([
      { name: field_name, values: ReadOnlyArray::from_array(values) },
    ])
  }
}

///|
fn[D : @lib.Directory] search_index_with_highlights(
  directory : D,
  field_name : String,
  query_text : String,
  limit : Int,
  mode : QueryMode,
  fragment_size : Int,
) -> SearchReport raise CliError {
  guard fragment_size > 0 else {
    fail("--fragment-size must be a positive integer")
  }
  guard mode != QueryString else {
    fail("--highlight is not yet supported with --query-string")
  }
  search_index_internal(
    directory,
    field_name,
    query_text,
    limit,
    mode,
    Some(fragment_size),
  )
}

///|
fn[D : @lib.Directory] search_index_internal(
  directory : D,
  field_name : String,
  query_text : String,
  limit : Int,
  mode : QueryMode,
  fragment_size : Int?,
) -> SearchReport raise CliError {
  guard limit > 0 else { fail("--limit must be a positive integer") }
  let reader = @lib.IndexReader::open(directory) catch {
    error => fail(error.to_string())
  }
  let schema_option = reader.schema() catch { error => fail(error.to_string()) }
  let schema = match schema_option {
    Some(schema) => schema
    None => fail("index does not contain a schema")
  }
  let field_id = match schema.field(field_name) {
    Some(field_id) => field_id
    None => fail("unknown field: \{field_name}")
  }
  guard schema.is_indexed(field_id) else {
    fail("field is not indexed: \{field_name}")
  }
  let parser = @lib.QueryParser::new(
    schema,
    @lib.TokenizerManager::with_defaults(),
  )
  let highlighter = @lib.Highlighter::new(
    schema,
    @lib.TokenizerManager::with_defaults(),
  )
  let query : &@lib.Query = match mode {
    Or =>
      parser.parse_query(field_id, query_text) catch {
        error => fail(error.to_string())
      }
    And => {
      parser.set_conjunction_by_default(true)
      parser.parse_query(field_id, query_text) catch {
        error => fail(error.to_string())
      }
    }
    Phrase =>
      parser.parse_phrase(field_id, query_text) catch {
        error => fail(error.to_string())
      }
    QueryString => {
      let query_string_parser = @lib.QueryStringParser::new(
        schema,
        @lib.TokenizerManager::with_defaults(),
        [field_id],
      )
      query_string_parser.parse_query(query_text) catch {
        error => fail(error.to_string())
      }
    }
  }
  let searcher = reader.searcher() catch { error => fail(error.to_string()) }
  let hits : Array[SearchHitReport] = []
  for hit in searcher.search(query, @lib.TopKCollector::new(limit)) {
    let (fields, highlights) = match searcher.doc(hit.address) {
      Some(document) =>
        (
          stored_fields(schema, document),
          match fragment_size {
            Some(size) =>
              highlight_field(
                highlighter, field_id, field_name, query_text, mode, size, document,
              )
            None => ReadOnlyArray::from_array([])
          },
        )
      None => (ReadOnlyArray::from_array([]), ReadOnlyArray::from_array([]))
    }
    hits.push({
      segment_ord: hit.address.segment_ord,
      doc_id: hit.address.doc_id.value,
      score: hit.score,
      fields,
      highlights,
    })
  }
  reader.close()
  {
    field: field_name,
    query: query_text,
    hits: ReadOnlyArray::from_array(hits),
  }
}

///|
fn[D : @lib.Directory] inspect_index(
  directory : D,
) -> InspectReport raise CliError {
  let reader = @lib.IndexReader::open(directory) catch {
    error => fail(error.to_string())
  }
  let schema_option = reader.schema() catch { error => fail(error.to_string()) }
  let schema = match schema_option {
    Some(schema) => schema
    None => fail("index does not contain a schema")
  }
  let fields : Array[FieldReport] = []
  for index in 0.. fail(error.to_string()) }
  let report = {
    generation: reader.generation() catch {
      error => fail(error.to_string())
    },
    segment_count: reader.segment_count() catch {
      error => fail(error.to_string())
    },
    document_count: searcher.doc_count(),
    fields: ReadOnlyArray::from_array(fields),
  }
  reader.close()
  report
}

///|
fn stored_fields_json(fields : ReadOnlyArray[StoredFieldReport]) -> Json {
  let object : Map[String, Json] = Map([])
  for field in fields {
    let values : Array[Json] = []
    for value in field.values {
      values.push(Json::string(value))
    }
    object[field.name] = Json::array(values)
  }
  Json::object(object)
}

///|
fn highlight_fragment_json(value : HighlightValueReport) -> Json {
  let matches : Array[Json] = []
  for span in value.fragment.matches {
    matches.push(
      Json::object({
        "start_offset": Json::number(span.start_offset.to_double()),
        "end_offset": Json::number(span.end_offset.to_double()),
      }),
    )
  }
  Json::object({
    "value_index": Json::number(value.value_index.to_double()),
    "text": Json::string(value.fragment.text),
    "start_offset": Json::number(value.fragment.start_offset.to_double()),
    "end_offset": Json::number(value.fragment.end_offset.to_double()),
    "prefix_omitted": Json::boolean(value.fragment.prefix_omitted),
    "suffix_omitted": Json::boolean(value.fragment.suffix_omitted),
    "matches": Json::array(matches),
  })
}

///|
fn highlights_json(fields : ReadOnlyArray[HighlightFieldReport]) -> Json {
  let object : Map[String, Json] = Map([])
  for field in fields {
    let values : Array[Json] = []
    for value in field.values {
      values.push(highlight_fragment_json(value))
    }
    object[field.name] = Json::array(values)
  }
  Json::object(object)
}

///|
fn rendered_highlights_json(
  fields : ReadOnlyArray[HighlightFieldReport],
) -> Json {
  let object : Map[String, Json] = Map([])
  for field in fields {
    let values : Array[Json] = []
    for value in field.values {
      let fragment = value.fragment
      let mut rendered = fragment.render("[[", "]]")
      if fragment.prefix_omitted {
        rendered = "…" + rendered
      }
      if fragment.suffix_omitted {
        rendered += "…"
      }
      values.push(
        Json::object({
          "value_index": Json::number(value.value_index.to_double()),
          "text": Json::string(rendered),
        }),
      )
    }
    object[field.name] = Json::array(values)
  }
  Json::object(object)
}

///|
fn render_index_report(report : IndexReport, format : OutputFormat) -> String {
  match format {
    Text =>
      "indexed \{report.document_count} documents; generation=\{report.generation} segments=\{report.segment_count}"
    Json =>
      Json::object({
        "documents_indexed": Json::number(report.document_count.to_double()),
        "generation": Json::number(report.generation.to_double()),
        "segments": Json::number(report.segment_count.to_double()),
      }).stringify()
  }
}

///|
fn render_search_report(report : SearchReport, format : OutputFormat) -> String {
  match format {
    Json => {
      let hits : Array[Json] = []
      for hit in report.hits {
        let object : Map[String, Json] = {
          "segment": Json::number(hit.segment_ord.to_double()),
          "doc": Json::number(hit.doc_id.to_double()),
          "score": Json::number(hit.score),
          "fields": stored_fields_json(hit.fields),
        }
        if hit.highlights.length() > 0 {
          object["highlights"] = highlights_json(hit.highlights)
        }
        hits.push(Json::object(object))
      }
      Json::object({
        "field": Json::string(report.field),
        "query": Json::string(report.query),
        "hits": Json::array(hits),
      }).stringify()
    }
    Text => {
      if report.hits.length() == 0 {
        return "no hits"
      }
      let mut output = ""
      for index in 0.. 0 {
          output += "\n"
        }
        output += "\{index + 1}. score=\{hit.score} segment=\{hit.segment_ord} doc=\{hit.doc_id} fields=\{stored_fields_json(hit.fields).stringify()}"
        if hit.highlights.length() > 0 {
          output += " highlights=\{rendered_highlights_json(hit.highlights).stringify()}"
        }
      }
      output
    }
  }
}

///|
fn render_inspect_report(
  report : InspectReport,
  format : OutputFormat,
) -> String {
  match format {
    Json => {
      let fields : Array[Json] = []
      for field in report.fields {
        fields.push(
          Json::object({
            "id": Json::number(field.id.to_double()),
            "name": Json::string(field.name),
            "indexed": Json::boolean(field.indexed),
            "stored": Json::boolean(field.stored),
            "tokenizer": Json::string(field.tokenizer),
          }),
        )
      }
      Json::object({
        "generation": Json::number(report.generation.to_double()),
        "segments": Json::number(report.segment_count.to_double()),
        "documents": Json::number(report.document_count.to_double()),
        "fields": Json::array(fields),
      }).stringify()
    }
    Text => {
      let mut output = "generation=\{report.generation} segments=\{report.segment_count} documents=\{report.document_count}"
      for field in report.fields {
        output += "\nfield \{field.id}: \{field.name} indexed=\{field.indexed} stored=\{field.stored} tokenizer=\{field.tokenizer}"
      }
      output
    }
  }
}

///|
fn common_output_args() -> (
  Array[@argparse.FlagArg],
  Array[@argparse.OptionArg],
) {
  (
    [
      @argparse.FlagArg(
        "json",
        long="json",
        about="Emit JSON output (shorthand for --format json)",
      ),
    ],
    [
      @argparse.OptionArg(
        "format",
        long="format",
        about="Output format: text or json",
        default_values=["text"],
      ),
    ],
  )
}

///|
fn cli_command() -> @argparse.Command {
  let (output_flags, output_options) = common_output_args()
  @argparse.Command(
    "moonsearch",
    about="Index and search UTF-8 JSONL documents with MoonSearch",
    version="0.9.0",
    subcommand_required=true,
    arg_required_else_help=true,
    subcommands=[
      @argparse.Command(
        "index",
        about="Append JSONL documents as a new immutable segment",
        flags=output_flags,
        options=[
          @argparse.OptionArg(
            "index",
            long="index",
            about="Index directory",
            required=true,
          ),
          @argparse.OptionArg(
            "input",
            long="input",
            about="UTF-8 JSONL input file",
            required=true,
          ),
          @argparse.OptionArg(
            "field",
            long="field",
            about="Indexed and stored text field (repeatable)",
            action=@argparse.OptionAction::Append,
            required=true,
          ),
          @argparse.OptionArg(
            "tokenizer",
            long="tokenizer",
            about="Built-in tokenizer: default, whitespace, raw, or en_stem",
            default_values=["default"],
          ),
          ..output_options,
        ],
      ),
      @argparse.Command(
        "search",
        about="Search an index with raw field text or Query-string syntax",
        flags=[
          @argparse.FlagArg(
            "phrase",
            long="phrase",
            about="Run an exact phrase query",
          ),
          @argparse.FlagArg(
            "highlight",
            long="highlight",
            about="Include UTF-8-safe highlighted fragments",
          ),
          @argparse.FlagArg(
            "query-string",
            long="query-string",
            about="Parse --query as field-aware Query-string syntax",
          ),
          ..output_flags,
        ],
        options=[
          @argparse.OptionArg(
            "index",
            long="index",
            about="Index directory",
            required=true,
          ),
          @argparse.OptionArg(
            "field",
            long="field",
            about="Indexed field name (default field for Query-string)",
            required=true,
          ),
          @argparse.OptionArg(
            "query",
            long="query",
            about="Raw query text or Query-string expression",
            required=true,
          ),
          @argparse.OptionArg(
            "limit",
            long="limit",
            about="Maximum number of hits",
            default_values=["10"],
          ),
          @argparse.OptionArg(
            "operator",
            long="operator",
            about="Multi-term operator: or or and",
            default_values=["or"],
          ),
          @argparse.OptionArg(
            "fragment-size",
            long="fragment-size",
            about="Maximum Unicode characters per highlight fragment",
            default_values=["160"],
          ),
          ..output_options,
        ],
      ),
      @argparse.Command(
        "inspect",
        about="Show index snapshot and schema metadata",
        flags=output_flags,
        options=[
          @argparse.OptionArg(
            "index",
            long="index",
            about="Index directory",
            required=true,
          ),
          ..output_options,
        ],
      ),
    ],
  )
}

///|
fn value(matches : @argparse.Matches, name : String) -> String raise CliError {
  match matches.values.get(name) {
    Some(values) if values.length() > 0 => values[0]
    _ => fail("missing option --\{name}")
  }
}

///|
fn output_format(matches : @argparse.Matches) -> OutputFormat raise CliError {
  if matches.flags.get("json") is Some(true) {
    return Json
  }
  match value(matches, "format") {
    "text" => Text
    "json" => Json
    format => fail("unsupported output format: \{format}")
  }
}

///|
fn query_mode(matches : @argparse.Matches) -> QueryMode raise CliError {
  if matches.flags.get("query-string") is Some(true) {
    guard !(matches.flags.get("phrase") is Some(true)) else {
      fail("--query-string cannot be combined with --phrase")
    }
    return QueryString
  }
  if matches.flags.get("phrase") is Some(true) {
    return Phrase
  }
  match value(matches, "operator") {
    "or" => Or
    "and" => And
    operator => fail("unsupported query operator: \{operator}")
  }
}

///|
fn execute_matches(matches : @argparse.Matches) -> String raise CliError {
  match matches.subcommand {
    Some(("index", child)) => {
      let index_path = value(child, "index")
      let input_path = value(child, "input")
      let input = @fs.read_file_to_string(input_path, encoding="utf8") catch {
        @fs.IOError::IOError(message) => fail("cannot read input: \{message}")
      }
      let fields = match child.values.get("field") {
        Some(fields) => fields
        None => fail("at least one --field is required")
      }
      let directory = @lib.FsDirectory::new(index_path) catch {
        error => fail(error.to_string())
      }
      render_index_report(
        index_jsonl(directory, input, fields, value(child, "tokenizer")),
        output_format(child),
      )
    }
    Some(("search", child)) => {
      let limit = @string.parse_int(value(child, "limit")) catch {
        _ => fail("--limit must be a positive integer")
      }
      let fragment_size = @string.parse_int(value(child, "fragment-size")) catch {
        _ => fail("--fragment-size must be a positive integer")
      }
      guard fragment_size > 0 else {
        fail("--fragment-size must be a positive integer")
      }
      let directory = @lib.FsDirectory::new(value(child, "index")) catch {
        error => fail(error.to_string())
      }
      render_search_report(
        if child.flags.get("highlight") is Some(true) {
          search_index_with_highlights(
            directory,
            value(child, "field"),
            value(child, "query"),
            limit,
            query_mode(child),
            fragment_size,
          )
        } else {
          search_index(
            directory,
            value(child, "field"),
            value(child, "query"),
            limit,
            query_mode(child),
          )
        },
        output_format(child),
      )
    }
    Some(("inspect", child)) => {
      let directory = @lib.FsDirectory::new(value(child, "index")) catch {
        error => fail(error.to_string())
      }
      render_inspect_report(inspect_index(directory), output_format(child))
    }
    Some((name, _)) => fail("unsupported subcommand: \{name}")
    None => fail("a subcommand is required")
  }
}