squirrel

import argv
import envoy
import filepath
import glam/doc.{type Document}
import gleam/bool
import gleam/dict.{type Dict}
import gleam/int
import gleam/io
import gleam/list
import gleam/option.{type Option, None, Some}
import gleam/result
import gleam/string
import gleam/uri.{Uri}
import gleam_community/ansi
import glexer
import glexer/token
import simplifile
import squirrel/internal/database/postgres
import squirrel/internal/error.{
  type Error, CannotOverwriteExistingFile, CannotReadFile, CannotWriteToFile,
  InvalidConnectionString, OutdatedFile,
}
import squirrel/internal/project
import squirrel/internal/query.{type TypedQuery}
import term_size

const squirrel_version = "v4.7.0"

/// 🐿️ Performs code generation for your Gleam project.
///
/// `squirrel` is not configurable and will discover the queries to generate
/// code for by relying on a conventional project's structure:
/// - `squirrel` first looks for all directories called `sql` under the `src`
///   directory of your Gleam project, and reads all the `*.sql` files in there
///   (in glob terms `src/**/sql/*.sql`).
/// - Each `*.sql` file _must contain a single query_ as it is turned into a
///   Gleam function with the same name.
/// - All functions coming from the same `sql` directory will be grouped under
///   a Gleam file called `sql.gleam` at the same level: given a `src/$PATH/sql`
///   directory, you'll end up with a generated `src/$PATH/sql.gleam` file.
///
/// > ⚠️ In order to generate type safe code, `squirrel` has to connect
/// > to your Postgres database. To know what host, user, etc. values to use
/// > when connecting, it will read the `DATABASE_URL` env variable that has to
/// > be a valid connection string with the following format:
/// >
/// > ```txt
/// > postgres://user:password@host:port/database?connect_timeout=seconds
/// > ```
/// >
/// > If a `DATABASE_URL` variable is not set, Squirrel will instead read your
/// > [Postgres env variables](https://www.postgresql.org/docs/current/libpq-envars.html)
/// > and use the following defaults if one is not set:
/// > - `PGHOST`: `"localhost"`
/// > - `PGPORT`: `5432`
/// > - `PGUSER`: `"postgres"`
/// > - `PGDATABASE`: the name of your Gleam project
/// > - `PGPASSWORD`: `""`
/// > - `PGCONNECT_TIMEOUT`: `5` seconds
///
/// > ⚠️ The generated code relies on the
/// > [`pog`](https://hexdocs.pm/pog/) package to work, so make sure to add
/// > that dependency to your project.
///
pub fn main() {
  case parse_cli_args() {
    Error(Nil) -> {
      help_text()
      |> doc.to_string(term_width())
      |> io.println

      exit(1)
    }

    Ok(mode) ->
      case result.try(connection_options(), postgres.connect_and_authenticate) {
        Error(error) -> {
          error.to_doc(error)
          |> doc.to_string(term_width())
          |> io.println

          exit(1)
        }

        Ok(connection) -> {
          let generated_queries =
            walk(project.src())
            |> dict.merge(walk(project.test_()))
            |> dict.merge(walk(project.dev()))
            |> generate_queries(connection)

          let #(report, status_code) = case mode {
            GenerateCode ->
              case ensure_no_errors(generated_queries) {
                Error(errors) -> report_errors(errors)
                Ok(generated_queries) ->
                  generated_queries
                  |> write_queries
                  |> report_written_queries
              }

            CheckGeneratedCode ->
              generated_queries
              |> check_queries
              |> report_checked_queries
          }

          io.println(report)
          exit(status_code)
        }
      }
  }
}

fn ensure_no_errors(
  generated_queries: Dict(String, #(List(TypedQuery), List(Error))),
) -> Result(Dict(String, List(TypedQuery)), List(Error)) {
  let all_errors =
    dict.fold(generated_queries, [], fn(errors, _directory, queries_and_errors) {
      let #(_queries, new_errors) = queries_and_errors
      list.append(new_errors, errors)
    })

  case all_errors {
    [_, ..] -> Error(all_errors)
    [] -> Ok(dict.map_values(generated_queries, fn(_, queries) { queries.0 }))
  }
}

// --- CLI ARGS PARSING --------------------------------------------------------

type Mode {
  GenerateCode
  CheckGeneratedCode
}

fn parse_cli_args() -> Result(Mode, Nil) {
  case argv.load().arguments {
    [] -> Ok(GenerateCode)
    ["check"] -> Ok(CheckGeneratedCode)
    _ -> Error(Nil)
  }
}

fn help_text() -> Document {
  let nesting = 4
  let title_line = doc.from_string("🐿️  Squirrel " <> squirrel_version)

  let usage_line =
    doc.concat([
      doc.zero_width_string(ansi.yellow("Usage: ")),
      doc.zero_width_string(ansi.green("gleam run -m squirrel")),
      doc.from_string(" [COMMAND]"),
    ])

  let check_command =
    [
      doc.from_string("check  "),
      "checks the generated code is up to date with the sql files"
        |> flexible_string
        |> doc.nest(by: 7),
    ]
    |> doc.concat
    |> doc.group

  let commands =
    [
      doc.zero_width_string(ansi.yellow("Commands:")),
      doc.nest(check_command, by: nesting),
    ]
    |> doc.join(with: doc.nest(doc.line, by: nesting))

  [title_line, usage_line, commands]
  |> doc.join(with: doc.lines(2))
}

// --- CONNECTION OPTIONS PARSING ----------------------------------------------

/// Returns the connection options to use to connect to the database.
/// It first tries to read and parse a `DATABASE_URL` env variable (failing if
/// it has an invalid format).
///
/// If the `DATABASE_URL` variable is not set, it uses the Postgres env vars and
/// some defaults if any of those are not set.
///
fn connection_options() -> Result(postgres.ConnectionOptions, Error) {
  case envoy.get("DATABASE_URL") {
    Ok(url) ->
      parse_connection_url(url)
      |> result.replace_error(InvalidConnectionString(url))
    Error(_) -> Ok(connection_options_from_variables())
  }
}

const default_host = "localhost"

const default_user = "postgres"

const default_database = "database"

const default_password = ""

const default_port = 5432

const default_timeout = 5

/// Creates a `ConnectionOptions` reading values from env variables and falling
/// back to some defaults if any required one is not set.
///
fn connection_options_from_variables() -> postgres.ConnectionOptions {
  let host = envoy.get("PGHOST") |> result.unwrap(default_host)
  let user = envoy.get("PGUSER") |> result.unwrap(default_user)
  let password = envoy.get("PGPASSWORD") |> result.unwrap(default_password)
  let database =
    envoy.get("PGDATABASE")
    |> result.or(project.name())
    |> result.unwrap(default_database)
  let port =
    envoy.get("PGPORT")
    |> result.try(int.parse)
    |> result.unwrap(default_port)
  let timeout_seconds =
    envoy.get("PGCONNECT_TIMEOUT")
    |> result.try(int.parse)
    |> result.unwrap(default_timeout)

  postgres.ConnectionOptions(
    host:,
    port:,
    user:,
    password:,
    database:,
    timeout_seconds:,
  )
}

/// Parses a connection string into a `ConnectionOptions` failing if it has an
/// invalid format instead of silently producing a default one.
///
fn parse_connection_url(
  raw: String,
) -> Result(postgres.ConnectionOptions, Nil) {
  use uri <- result.try(uri.parse(raw))
  let Uri(scheme:, userinfo:, host:, port:, path:, query:, ..) = uri

  use parameters <- result.try(case query {
    None -> Ok([])
    Some(parameters) -> uri.parse_query(parameters)
  })
  use _ <- result.try(check_scheme(scheme))
  let #(user, password) = parse_user_and_password_from_userinfo(userinfo)
  let database = parse_database_from_path(path)
  use timeout <- result.try(case list.key_find(parameters, "connect_timeout") {
    Error(_) -> Ok(default_timeout)
    Ok(timeout) -> int.parse(timeout)
  })

  Ok(postgres.ConnectionOptions(
    host: host |> option.unwrap(default_host),
    port: port |> option.unwrap(default_port),
    user: user |> option.unwrap(default_user),
    password: password |> option.unwrap(default_password),
    database: database |> option.unwrap(default_database),
    timeout_seconds: timeout,
  ))
}

fn check_scheme(scheme: Option(String)) -> Result(Nil, Nil) {
  case scheme {
    Some("postgres") | Some("postgresql") | None -> Ok(Nil)
    Some(_) -> Error(Nil)
  }
}

fn parse_user_and_password_from_userinfo(
  userinfo: Option(String),
) -> #(Option(String), Option(String)) {
  case userinfo {
    None -> #(None, None)
    Some(userinfo) ->
      case string.split(userinfo, on: ":") {
        [user] -> #(Some(user), None)
        [user, password, ..] -> #(Some(user), Some(password))
        _ -> #(None, None)
      }
  }
}

fn parse_database_from_path(path: String) -> Option(String) {
  case string.split(path, "/") {
    ["", database, ..] -> Some(database)
    _ -> None
  }
}

/// Finds all `from/**/sql` directories and lists the full paths of the `*.sql`
/// files inside each one.
///
fn walk(from: String) -> Dict(String, List(String)) {
  case filepath.base_name(from) {
    "sql" -> {
      let assert Ok(files) = simplifile.read_directory(from)
      let files = {
        use file <- list.filter_map(files)
        use extension <- result.try(filepath.extension(file))
        use <- bool.guard(when: extension != "sql", return: Error(Nil))
        let file_name = filepath.join(from, file)
        case simplifile.is_file(file_name) {
          Ok(True) -> Ok(file_name)
          Ok(False) | Error(_) -> Error(Nil)
        }
      }
      dict.from_list([#(from, files)])
    }

    _ -> {
      let files = case simplifile.read_directory(from) {
        Ok(files) -> files
        Error(simplifile.Enoent) -> []
        Error(_) -> panic as { "couldn't read directory: " <> from }
      }
      let directories = {
        use file <- list.filter_map(files)
        let file_name = filepath.join(from, file)
        case simplifile.is_directory(file_name) {
          Ok(True) -> Ok(file_name)
          Ok(False) | Error(_) -> Error(Nil)
        }
      }

      list.map(directories, walk)
      |> list.fold(from: dict.new(), with: dict.merge)
    }
  }
}

// --- MODE: GENERATE QUERIES --------------------------------------------------

/// Given a dict of directories and their `*.sql` files, performs code
/// generation for each one, bundling all `*.sql` files under the same directory
/// into a single Gleam module.
///
fn generate_queries(
  directories: Dict(String, List(String)),
  connection: postgres.Context,
) -> Dict(String, #(List(TypedQuery), List(Error))) {
  use _directory, files <- dict.map_values(directories)

  let #(queries, errors) =
    list.map(files, query.from_file)
    |> result.partition

  case postgres.main(queries, connection) {
    Error(error) -> #([], [error, ..errors])
    Ok(#(queries, type_errors)) -> #(queries, list.append(errors, type_errors))
  }
}

/// Given the queries generated by `generate_queries`, tries to write those to
/// their own file and returns a dictionary that - for each file - holds the
/// number of queries that could be generated or the error that occurred trying
/// to write the query.
///
fn write_queries(
  queries: Dict(String, List(TypedQuery)),
) -> Dict(String, Result(Int, Error)) {
  use directory, queries <- dict.map_values(queries)
  let output_file = directory_to_output_file(directory)
  write_queries_to_file(queries, from: directory, to: output_file)
}

fn write_queries_to_file(
  queries: List(TypedQuery),
  from queries_directory: String,
  to file: String,
) -> Result(Int, Error) {
  use <- bool.guard(when: queries == [], return: Ok(0))
  let directory = filepath.directory_name(file)
  let _ = simplifile.create_directory_all(directory)

  let code =
    query.generate_code(squirrel_version, for: queries, from: queries_directory)

  safely_overwrite(to: file, content: code)
  |> result.replace(list.length(queries))
}

fn safely_overwrite(
  to file: String,
  content code: String,
) -> Result(Nil, Error) {
  case simplifile.read(file) |> result.map(classify_file_content) {
    // If the file is generated, empty, or doesn't exist at all then it's safe
    // to overwrite it!
    Ok(LikelyGenerated) | Ok(Empty) | Error(simplifile.Enoent) ->
      simplifile.write(code, to: file)
      |> result.map_error(CannotWriteToFile(file, _))

    // Otherwise, we can't safely replace its content.
    Ok(NotGenerated) -> Error(CannotOverwriteExistingFile(file))
    Error(reason) -> Error(CannotWriteToFile(file, reason))
  }
}

@internal
pub fn classify_file_content(content: String) -> FileOrigin {
  let likely_generated =
    // In newer versions of squirrel this is always at the beginning of the
    // file and it would be enough to check for this comment to establish if
    // a file is generated or not...
    string.contains(
      content,
      "> 🐿️ This module was generated automatically using",
    )
    // ...but in older versions that module comment is not present! So we
    // need to check if there's any function generated by squirrel.
    || string.contains(
      content,
      "> 🐿️ This function was generated automatically using",
    )

  case likely_generated {
    True -> LikelyGenerated
    False ->
      case string.trim(content) {
        "" -> Empty
        _ -> NotGenerated
      }
  }
}

@internal
pub type FileOrigin {
  /// The file was most likely generated by Squirrel. It can safely be
  /// overwritten as we do not risk deleting code written by a human.
  ///
  LikelyGenerated

  /// The file wasn't likely generated but it's only made of whitespace, so it's
  /// safe to override.
  Empty

  /// The file was most likely not generated by Squirrel. It is not safe to
  /// overwrite it, or we'd risk deleting code written by a human!
  NotGenerated
}

fn directory_to_output_file(directory: String) -> String {
  filepath.directory_name(directory)
  |> filepath.join("sql.gleam")
}

// --- MODE: CHECK GENERATED CODE ----------------------------------------------

fn check_queries(
  queries: Dict(String, #(List(query.TypedQuery), List(Error))),
) -> Dict(String, Result(Nil, List(Error))) {
  use directory, #(queries, errors) <- dict.map_values(queries)
  case errors {
    // If there's any error in the query generation (for example there's a piece
    // of sql code with a syntax error, or a column doesn't exist, ...)
    // we don't even try to check that the files match, but just report the
    // errors that took place.
    [_, ..] -> Error(errors)
    [] -> {
      let output_file = directory_to_output_file(directory)
      case simplifile.read(output_file) {
        Error(reason) -> Error([CannotReadFile(file: output_file, reason:)])
        Ok(actual_code) ->
          case check_queries_code(queries, actual_code) {
            Different -> Error([OutdatedFile(file: output_file)])
            Same -> Ok(Nil)
          }
      }
    }
  }
}

@internal
pub type CheckResult {
  Different
  Same
}

/// Checks that the given `file`'s code is the same that would be generated from
/// the given `queries`.
/// Returns `Same` if the code is ok, `Different` otherwise.
///
/// > ⚠️ Note how this comparison doesn't take comments and whitespace into
/// > account! If the only thing that changed are comments and/or formatting
/// > two files will still be considered the same.
///
fn check_queries_code(
  queries: List(TypedQuery),
  actual_code: String,
) -> CheckResult {
  // The code is compared ignoring any comments, so we have no need to know
  // what actual directory the queries come from: that information is only used
  // to generate better comments!
  let expected_code =
    query.generate_code(squirrel_version, for: queries, from: "check-queries")
  compare_code_snippets(actual_code, expected_code)
}

@internal
pub fn compare_code_snippets(
  actual_code: String,
  expected_code: String,
) -> CheckResult {
  let actual_tokens =
    glexer.new(actual_code)
    |> glexer.discard_comments
    |> glexer.discard_whitespace
    |> glexer.lex

  let expected_tokens =
    glexer.new(expected_code)
    |> glexer.discard_comments
    |> glexer.discard_whitespace
    |> glexer.lex

  compare_token_lists(expected_tokens, actual_tokens)
}

fn compare_token_lists(
  expected_tokens: List(#(token.Token, glexer.Position)),
  actual_tokens: List(#(token.Token, glexer.Position)),
) {
  case expected_tokens, actual_tokens {
    [], [] -> Same
    [], _ | _, [] -> Different
    [#(expected_token, _position), ..expected_rest],
      [#(actual_token, _position), ..actual_rest]
    ->
      case expected_token == actual_token {
        True -> compare_token_lists(expected_rest, actual_rest)
        False -> Different
      }
  }
}

// --- PRETTY REPORT PRINTING --------------------------------------------------

fn term_width() -> Int {
  term_size.columns()
  |> result.unwrap(80)
  |> int.min(80)
}

fn report_written_queries(
  directories: Dict(String, Result(Int, Error)),
) -> #(String, Int) {
  let #(ok, errors) = {
    use acc, _, outcome <- dict.fold(directories, #(0, []))
    let #(generated_queries, errors) = acc
    case outcome {
      Error(error) -> #(generated_queries, [error, ..errors])
      Ok(new_count) -> #(generated_queries + new_count, errors)
    }
  }

  let status_code = case errors {
    [_, ..] -> 1
    [] -> 0
  }

  let report = case ok, errors {
    0, [_, ..] ->
      errors_to_doc(errors)
      |> doc.to_string(term_width())

    0, [] ->
      [
        text_with_header(
          "🐿️  ",
          "I couldn't find any `*.sql` file to generate queries from.",
        ),
        doc.lines(2),
        flexible_string(
          "Hint: I look for all `*.sql` files in any directory called `sql`
under your project's `src`, `test`, and `dev` directories.",
        ),
      ]
      |> doc.concat
      |> doc.to_string(term_width())
      |> ansi.yellow

    n, [] ->
      text_with_header(
        "🐿️  ",
        "Generated "
          <> int.to_string(n)
          <> " "
          <> pluralise(n, "query", "queries")
          <> "!",
      )
      |> doc.to_string(term_width())
      |> ansi.green

    n, [_, ..] ->
      [
        errors_to_doc(errors),
        doc.lines(2),
        text_with_header(
          "🥜 ",
          "I could still generate "
            <> int.to_string(n)
            <> " "
            <> pluralise(n, "query", "queries")
            <> ".",
        ),
      ]
      |> doc.concat
      |> doc.to_string(term_width())
  }

  #(report, status_code)
}

fn errors_to_doc(errors: List(Error)) -> Document {
  list.map(errors, error.to_doc)
  |> doc.join(with: doc.lines(2))
}

fn report_errors(errors: List(Error)) -> #(String, Int) {
  let report = errors_to_doc(errors) |> doc.to_string(term_width())
  #(report, 1)
}

fn report_checked_queries(
  dirs: Dict(String, Result(Nil, List(Error))),
) -> #(String, Int) {
  let errors = {
    use all_errors, _, result <- dict.fold(dirs, [])
    case result {
      Error(errors) -> errors |> list.append(all_errors)
      Ok(_) -> all_errors
    }
  }

  let status_code = case errors {
    [_, ..] -> 1
    [] -> 0
  }

  let report = case errors {
    [_, ..] ->
      list.map(errors, error.to_doc)
      |> doc.join(with: doc.lines(2))
      |> doc.to_string(term_width())

    [] ->
      text_with_header("🐿️  ", "All good!")
      |> doc.to_string(term_width())
      |> ansi.green
  }

  #(report, status_code)
}

fn text_with_header(header: String, text: String) {
  [
    doc.from_string(header),
    flexible_string(text)
      |> doc.nest(by: string.length(header)),
  ]
  |> doc.concat
  |> doc.group
}

fn pluralise(count: Int, singular: String, plural: String) -> String {
  case count {
    1 -> singular
    _ -> plural
  }
}

fn flexible_string(string: String) -> Document {
  string.split(string, on: "\n")
  |> list.flat_map(string.split(_, on: " "))
  |> list.map(doc.from_string)
  |> doc.join(with: doc.flex_space)
  |> doc.group
}

// --- FFI ---------------------------------------------------------------------

@external(erlang, "squirrel_ffi", "exit")
fn exit(n: Int) -> Nil