eagle0/src/main/scala/net/eagle0/common/TsvUtils.scala

package net.eagle0.common

import java.io.{File, PrintWriter}
import java.net.URL

import scala.collection.immutable.ListMap
import scala.io.Source
import scala.util.{Try, Using}

object TsvUtils {
  // Loads a TSV where the first line are keys, and the rest of the column is a sequence of values
  private def loadColumnMaps(
      source: Source
  ): Map[String, Vector[String]] =
    MoreSeq
      .irregularTranspose(loadLines(source))
      .map(_.filterNot(_.isBlank))
      .map {
        case h +: t => h -> t
        case _      =>
          throw new IllegalArgumentException("invalid TSV for column maps")
      }
      .toMap

  def loadColumnMaps(url: URL): Try[Map[String, Vector[String]]] =
    Using(Source.fromURL(url, "UTF-8")) {
      loadColumnMaps
    }

  // Loads a tsv and returns a Vector of arrays of strings.
  def loadColumnMaps(path: String): Try[Map[String, Vector[String]]] =
    Using(Source.fromFile(path, "UTF-8")) {
      loadColumnMaps
    }

  def loadLines(url: URL): Try[Vector[Vector[String]]] =
    Using(Source.fromURL(url, "UTF-8"))(loadLines)

  // Loads a tsv and returns a Vector of arrays of strings.
  def loadLines(path: String): Try[Vector[Vector[String]]] =
    Using(Source.fromFile(path, "UTF-8"))(loadLines)

  def loadMaps(url: URL): Try[Vector[ListMap[String, String]]] =
    loadLines(url).map(linesToMaps)

  // Loads a tsv where the first line consists of headers, and returns an Vector of key-value maps from header->line value.
  def loadMaps(path: String): Try[Vector[ListMap[String, String]]] =
    loadLines(path).map(linesToMaps)

  def saveMaps(maps: Vector[Map[String, String]], url: URL): Unit =
    saveLines(mapsToLines(maps), url)

  private def loadLines(source: Source): Vector[Vector[String]] =
    source
      .getLines()
      .filterNot(_.startsWith("//"))
      .map(
        _.split("\t")
          .map(_.trim)
          .toVector
      )
      .toVector

  private def linesToMaps(
      lines: Vector[Vector[String]]
  ): Vector[ListMap[String, String]] =
    lines match {
      case Vector(_) => Vector.empty
      case h +: t    => t.map(arr => ToListMap(h.zip(arr)))
      case _         => Vector.empty
    }

  private def ToListMap(
      elts: Iterable[(String, String)]
  ): ListMap[String, String] = elts.foldLeft(ListMap[String, String]()) {
    case (lm, pair) =>
      lm + pair
  }

  private def mapsToLines(
      maps: Vector[Map[String, String]]
  ): Vector[Vector[String]] = {
    val headers    = maps.head.keys.toVector
    val otherLines = maps.map(oneMap => headers.map(oneMap))

    headers +: otherLines
  }

  private def saveLines(lines: Vector[Vector[String]], url: URL): Unit = {
    val file = new File(url.getPath)
    val pw   = new PrintWriter(file)

    try
      lines.foreach(line => pw.write(line.mkString("\t") + "\n"))
    finally
      pw.close()
  }
}