acrostic

import acrostic/helper
import acrostic/parser
import gleam/int
import gleam/io
import gleam/list
import gleam/option
import gleam/regex
import gleam/result
import gleam/string
import nibble
import nibble/lexer
import simplifile
import sprinkle.{format}

pub type Message {
  Message(id: Int, name: String, fields: List(parser.PbMessageField))
}

pub type Flags {
  Flags(enum_to_int: Bool, int_to_enum: Bool)
}

pub fn gen(protos: List(String), to out_path: String, flags flags: Flags) {
  protos
  |> list.map(fn(filepath) {
    case simplifile.read(from: filepath) {
      Ok(content) -> content
      Error(e) -> panic as string.inspect(e)
    }
  })
  |> list.fold("", string.append)
  |> generate_proto(out_path, flags)

  io.println("done")
}

fn generate_proto(text: String, out_path: String, flags: Flags) {
  let #(lexer, message_parser, enum_parser) = parser.parser()
  let enums = get_enums(text, lexer, enum_parser)
  let structs =
    get_structs(text, lexer, message_parser)
    |> list.map(fn(a) {
      Message(id: a.0, name: { a.1 }.name, fields: { a.1 }.fields)
    })
  let messages =
    get_messages(text, lexer, message_parser)
    |> list.map(fn(a) {
      Message(id: a.0, name: { a.1 }.name, fields: { a.1 }.fields)
    })

  let assert Ok(_) =
    "
    import gleam/list
    import gleam/bit_array
    import acrostic/encoding.{type FieldEncoder, FieldEncoder}
    import acrostic/decoding
    "
    |> simplifile.write(to: out_path)

  write_enums(enums, out_path, flags)
  // write structs
  let _ = case list.length(structs) > 0 {
    True -> {
      let _ =
        simplifile.append(
          to: out_path,
          contents: "// struct start -----------------------------------\n",
        )
      write_structs(structs, out_path)
    }
    _ -> Nil
  }
  // write messages
  let _ = case list.length(messages) > 0 {
    True -> {
      let _ =
        simplifile.append(
          to: out_path,
          contents: "// messages start -----------------------------------\n",
        )
      let _ = write_messages(messages, out_path)
      Nil
    }
    _ -> Nil
  }

  helper.cmd("gleam format " <> out_path)
}

// pub type Message {
//   Ping(msg: String)
//   Pong(msg: String)
// }

// f.name <> ": " <> to_gleam_ty(f.ty, f.repeated)
fn write_messages(messages: List(Message), out_path: String) {
  let assert Ok(_) =
    simplifile.append(to: out_path, contents: "pub type Message {\n")
  let assert Ok(_) =
    messages
    |> list.map(message_to_string(_, fn(field) {
      field.name <> ": " <> to_gleam_ty(field.ty, field.repeated)
    }))
    |> list.fold("", string.append)
    |> simplifile.append(to: out_path)

  let assert Ok(_) = simplifile.append(to: out_path, contents: "}\n\n")

  // encoding gen
  let body =
    messages
    |> list.map(fn(msg) {
      format(
        "
          {message} -> {
        {fields}
          }
        ",
        [
          #("message", message_to_string(msg, fn(f) { f.name })),
          #(
            "fields",
            msg.fields
              |> list.map(get_field_encoding(_, ""))
              |> list.fold(
                "<<" <> int.to_string(msg.id) <> ":big-size(16)>>\n",
                string.append,
              ),
          ),
        ],
      )
    })
    |> list.fold("", string.append)

  let assert Ok(_) =
    "
    pub fn encode(msg: Message) -> BitArray {
    case msg {
        {body}
    }
  }
  "
    |> format([#("body", body)])
    |> simplifile.append(to: out_path)

  // decoding
  let assert Ok(_) =
    messages
    |> list.map(fn(msg) { get_message_case_code(msg) })
    |> list.fold(
      "
        pub fn decode_message(msg: Message, binary: BitArray) -> Message {
        case binary {
          <<>> -> msg
        _ -> case msg {
      ",
      string.append,
    )
    |> string.append(
      "
            }
          }
        }
      ",
    )
    |> simplifile.append(to: out_path)

  // pub const hello_default = Hello(1,2,3)
  messages
  |> list.each(fn(msg) {
    let assert Ok(_) =
      "pub const {name}_default = {value}\n"
      |> format([
        #("name", pascal_to_snake(msg.name)),
        #(
          "value",
          msg.fields
            |> list.map(fn(f) { get_default_value(f.ty, f.repeated) })
            |> list.fold(msg.name <> "(", fn(a, b) { a <> b <> "," })
            |> string.append(")"),
        ),
      ])
      |> simplifile.append(to: out_path)
  })

  write_decode(messages, out_path)
}

fn write_decode(messages: List(Message), out_path: String) {
  let assert Ok(_) =
    messages
    |> list.map(fn(msg) {
      format("{id} -> decode_message({default},binary)\n", [
        #("id", int.to_string(msg.id)),
        #("default", pascal_to_snake(msg.name) <> "_default"),
      ])
    })
    |> list.fold(
      "
      pub fn decode(binary: BitArray) -> Message {
        let assert <<id:big-size(16), binary:bits>> = binary
        case id { 
    ",
      string.append,
    )
    |> string.append(
      "_ -> panic
        }
      }
    ",
    )
    |> simplifile.append(to: out_path)
}

//     Hello(session, texts) -> {
//       let #(key, bin) = decoding.decode_key(bin)
//       case key.field_number {
//         1 -> {
//           let #(session, bin) = decoding.decode_varint(bin)
//           decode_message(Hello(session, texts), bin)
//         }
//         _ -> todo
//       }
//     }
fn get_message_case_code(msg: Message) -> String {
  msg.fields
  |> list.map(fn(f) {
    format(
      "
        {field_number} -> {
          let #({field_name}, binary) = {reader}
          decode_message({message}, binary)
        }
      ",
      [
        #("reader", get_reader_string(f)),
        #("field_number", int.to_string(f.tag)),
        #("field_name", {
          case f.repeated, field_is_packed(f) {
            True, False -> "value"
            _, _ -> f.name
          }
        }),
        #(
          "message",
          message_to_string(msg, fn(f2) {
            case f.name == f2.name, f2.repeated, field_is_packed(f2) {
              True, True, False ->
                format("list.append({field_name}, [value])", [
                  #("field_name", f2.name),
                ])
              _, _, _ -> f2.name
            }
          }),
        ),
      ],
    )
  })
  |> list.fold(message_to_string(msg, fn(f) { f.name }) <> "
    -> {
      let #(key, binary) = decoding.read_key(binary)
      case key.field_number {
  ", string.append)
  |> string.append(
    "
          _ -> panic
      }
    }
  ",
  )
}

// pub type Item {
//   Item(id: Int, num: Int)
// }
fn write_structs(structs: List(Message), out_path: String) {
  structs
  |> list.each(fn(struct) {
    // define gen
    let assert Ok(_) =
      simplifile.append(
        to: out_path,
        contents: "pub type " <> struct.name <> " {\n",
      )
    let assert Ok(_) =
      simplifile.append(
        to: out_path,
        contents: message_to_string(struct, fn(field) {
          field.name <> ": " <> to_gleam_ty(field.ty, field.repeated)
        })
          <> "}\n\n",
      )

    // default
    // pub const item_default = Item(0, 0)
    let assert Ok(_) =
      "pub const {name}_default = {value}\n"
      |> format([
        #("name", pascal_to_snake(struct.name)),
        #(
          "value",
          struct.fields
            |> list.map(fn(f) { get_default_value(f.ty, f.repeated) })
            |> list.fold(struct.name <> "(", fn(a, b) { a <> b <> "," })
            |> string.append(")"),
        ),
      ])
      |> simplifile.append(to: out_path)
    // encoding gen
    // fn encode_item(item: Item) -> BitArray {
    //   <<>>
    //   |> encoding.encode_int_field(1, item.id)
    //   |> encoding.encode_int_field(2, item.num)
    // }
    let assert Ok(_) =
      format(
        "
          pub fn encode_{name}({name}: {type}) -> BitArray {
            {body}
          }
       ",
        [
          #("name", pascal_to_snake(struct.name)),
          #("type", struct.name),
          #(
            "body",
            struct.fields
              |> list.map(get_field_encoding(
                _,
                pascal_to_snake(struct.name) <> ".",
              ))
              |> list.fold("<<>>\n", string.append),
          ),
        ],
      )
      |> simplifile.append(to: out_path)
    // decoding gen
    let assert Ok(_) =
      format(
        "
      pub fn decode_to_{name}(binary: BitArray, {name}: {typename}) -> Result({typename}, String) {
        case binary {
          <<>> -> Ok({name})
          _ -> {
            let #(key, binary) = decoding.read_key(binary)
            case key.field_number {
              {body}
                _ -> panic
              }
            }
          }
        }
      ",
        [
          #("name", pascal_to_snake(struct.name)),
          #("typename", struct.name),
          #(
            "body",
            struct.fields
              |> list.map(fn(f) {
                format(
                  "
                  {field_number} -> {
                    let #({var_name}, binary) = {reader}
                    decode_to_{name}(binary, {typename}(..{name}, {field_name}: {field_value}))
                  }
                ",
                  [
                    #("name", pascal_to_snake(struct.name)),
                    #("typename", struct.name),
                    #("reader", get_reader_string(f)),
                    #("field_number", int.to_string(f.tag)),
                    #("field_name", f.name),
                    #("var_name", {
                      case f.repeated, field_is_packed(f) {
                        True, False -> "value"
                        _, _ -> f.name
                      }
                    }),
                    #("field_value", {
                      case f.repeated, field_is_packed(f) {
                        True, False ->
                          format("list.append({name}.{field_name}, [value])", [
                            #("name", pascal_to_snake(struct.name)),
                            #("field_name", f.name),
                          ])
                        _, _ -> f.name
                      }
                    }),
                  ],
                )
              })
              |> list.fold("", string.append),
          ),
        ],
      )
      |> simplifile.append(to: out_path)
    // read_xxxx gen
    // pub fn read_item(binary: BitArray) -> #(Item, BitArray) {
    //   decoding.read_len_field(binary, decode_to_item(_, item_default))
    // }
    let assert Ok(_) =
      format(
        "
        pub fn read_{name}(binary: BitArray) -> #({typename}, BitArray) {
          decoding.read_len_field(binary, decode_to_{name}(_, {name}_default))
        }
      ",
        [#("name", pascal_to_snake(struct.name)), #("typename", struct.name)],
      )
      |> simplifile.append(to: out_path)

    write_struct_field_encoder(struct, out_path)
  })
}

fn write_struct_field_encoder(struct: Message, out_path: String) {
  let assert Ok(_) =
    format(
      "
      pub const {name}_field_encoder = FieldEncoder(
        {wire_type},
        {encoder}
      )
    ",
      [
        #("name", pascal_to_snake(struct.name)),
        #("wire_type", "encoding.WireLenTy"),
        #("encoder", "encode_" <> pascal_to_snake(struct.name)),
      ],
    )
    |> simplifile.append(to: out_path)
  Nil
}

// fn to_gleam_ty(ty: String, repeated: Bool) -> String {
//   let ty = case ty {
//     // string
//     "string" -> "String"
//     // varint
//     "int32" | "int64" | "uint32" | "uint64" -> "Int"
//     "bool" -> "Bool"
//     // i64
//     "fixed64" | "sfixed64" | "double" -> "Float"
//     // i32
//     "fixed32" | "sfixed32" | "float" -> "Float"
//     // Custom Type (Struct | Enum)
//     x -> x
//   }
// decoding.decode_varint(binary)
fn get_reader_string(field: parser.PbMessageField) -> String {
  case field.repeated {
    True -> {
      case field.ty {
        "string" -> "decoding.read_string(binary)"
        "int32" | "int64" | "uint32" | "uint64" ->
          "decoding.read_len_packed_field(binary, decoding.read_varint)"
        "bool" -> "decoding.read_len_packed_field(binary, decoding.read_bool)"
        "fixed64" | "sfixed64" | "double" ->
          "decoding.read_len_packed_field(binary, decoding.read_i64)"
        "fixed32" | "sfixed32" | "float" ->
          "decoding.read_len_packed_field(binary, decoding.read_i32)"
        // Custom Type (Enum)
        ty -> "read_" <> pascal_to_snake(ty) <> "(binary)"
      }
    }
    False -> {
      case field.ty {
        "string" -> "decoding.read_string(binary)"
        "int32" | "int64" | "uint32" | "uint64" ->
          "decoding.read_varint(binary)"
        "bool" -> "decoding.read_bool(binary)"
        "fixed64" | "sfixed64" | "double" -> "decoding.read_i64(binary)"
        "fixed32" | "sfixed32" | "float" -> "decoding.read_i32(binary)"
        // Custom Type (Enum)
        ty -> "read_" <> pascal_to_snake(ty) <> "(binary)"
      }
    }
  }
}

fn field_is_packed(field: parser.PbMessageField) -> Bool {
  case field.ty {
    "int32" | "int64" | "uint32" | "uint64" if field.repeated -> True
    "bool" if field.repeated -> True
    "fixed64" | "sfixed64" | "double" if field.repeated -> True
    "fixed32" | "sfixed32" | "float" if field.repeated -> True
    // Custom Type (Enum)
    _ -> False
  }
}

// fn get_len_docoder_string(field: parser.Field) -> String {
//   case field.ty {
//     "string" -> "decoding.decode_to_string"
//     // Custom Type (Struct)
//     ty -> "decode_to_" <> ty
//   }
// }

fn get_default_value(ty: String, repeated: Bool) -> String {
  case to_gleam_ty(ty, False) {
    _any if repeated -> "[]"
    "Int" -> "0"
    "Bool" -> "False"
    "Float" -> "0.0"
    "String" -> "\"\""
    // Enum or Struct
    x -> pascal_to_snake(x) <> "_default"
  }
}

// |> bit_array.append(encoding.encode_field(1, 1, encoding.int_field_encoder))
fn get_field_encoding(
  field: parser.PbMessageField,
  field_prefix: String,
) -> String {
  format(
    "|> bit_array.append(encoding.{encode_field}({tag}, {value}, {field_encoder}))",
    [
      #("encode_field", case field.repeated {
        True -> "encode_repeated_field"
        False -> "encode_field"
      }),
      #("tag", int.to_string(field.tag)),
      #("value", field_prefix <> field.name),
      #("field_encoder", {
        case field.ty {
          // string
          "string" -> "encoding.string_field_encoder"
          // varint
          "int32" | "int64" | "uint32" | "uint64" ->
            "encoding.int_field_encoder"
          "bool" -> "encoding.bool_field_encoder"
          // i64
          "fixed64" | "sfixed64" | "double" -> "encoding.i64_field_encoder"
          // i32
          "fixed32" | "sfixed32" | "float" -> "encoding.i32_field_encoder"
          // Custom Type (Struct | Enum)
          x -> pascal_to_snake(x) <> "_field_encoder"
        }
      }),
    ],
  )
}

// fn float_wire_type(ty: String) -> String {
//   case ty {
//     "fixed32" | "sfixed32" | "float" -> "i32"
//     "fixed64" | "sfixed64" | "double" -> "i64"
//     x -> panic as { "Invalid float type: " <> x }
//   }
// }

// pub type Season {
//   Spring
//   Summer
// }
fn write_enums(enums: List(parser.PbEnum), out_path: String, flags: Flags) {
  enums
  |> list.each(fn(enum) {
    // define gen
    let assert Ok(_) =
      simplifile.append(
        to: out_path,
        contents: "pub type " <> enum.name <> " {\n",
      )
    enum.fields
    |> list.each(fn(field) {
      let assert Ok(_) =
        simplifile.append(to: out_path, contents: "  " <> field.name <> "\n")
    })
    let assert Ok(_) = simplifile.append(to: out_path, contents: "}\n\n")

    // default value
    // pub const user_status_default = Idle
    let assert Ok(_) =
      "pub const {name}_default = {value}\n"
      |> format([
        #("name", pascal_to_snake(enum.name)),
        #("value", case enum.fields |> list.first {
          Ok(f) -> f.name
          _ -> panic as { "Invalid empty enum: " <> enum.name }
        }),
      ])
      |> simplifile.append(to: out_path)

    // encoding gen
    // pub fn encode_item(item: Item) -> BitArray {
    //   ...
    // }
    let assert Ok(_) =
      format(
        "
          pub fn encode_{name}({name}: {type}) -> BitArray {
            {body}
          }
        ",
        [
          #("name", pascal_to_snake(enum.name)),
          #("type", enum.name),
          #(
            "body",
            enum.fields
              |> list.map(fn(f) {
                format("    {key} -> encoding.encode_varint({value})\n", [
                  #("key", f.name),
                  #("value", int.to_string(f.tag)),
                ])
              })
              |> list.fold(
                "case " <> pascal_to_snake(enum.name) <> " {\n",
                string.append,
              )
              |> string.append("  }"),
          ),
        ],
      )
      |> simplifile.append(to: out_path)
    // decoding gen

    let assert Ok(_) =
      format(
        "
          pub fn read_{name}(binary: BitArray) -> #({type}, BitArray) {
            let #(num, binary) = decoding.read_varint(binary)
            {body}
          }
        ",
        [
          #("name", pascal_to_snake(enum.name)),
          #("type", enum.name),
          #(
            "body",
            enum.fields
              |> list.map(fn(f) {
                format("    {key} -> {value}\n", [
                  #("key", int.to_string(f.tag)),
                  #("value", "#(" <> f.name <> ", binary)"),
                ])
              })
              |> list.fold("case num {\n", string.append)
              |> string.append("  _ -> panic\n")
              |> string.append("  }\n"),
          ),
        ],
      )
      |> simplifile.append(to: out_path)

    // enum_to_int
    let _ = case flags.enum_to_int {
      True -> write_enum_to_int(enum, out_path)
      _ -> Nil
    }

    // int_to_enum
    let _ = case flags.int_to_enum {
      True -> write_int_to_enum(enum, out_path)
      _ -> Nil
    }

    // field encoder
    write_enum_field_encoder(enum, out_path)
  })
}

// pub const i32_field_encoder = FieldEncoder(
//   wire_type: WireI32Ty,
//   encode: encode_i32,
// )
fn write_enum_field_encoder(enum: parser.PbEnum, out_path: String) {
  let assert Ok(_) =
    format(
      "
      pub const {name}_field_encoder = FieldEncoder(
        {wire_type},
        {encoder}
      )
    ",
      [
        #("name", pascal_to_snake(enum.name)),
        #("wire_type", "encoding.WireVarIntTy"),
        #("encoder", "encode_" <> pascal_to_snake(enum.name)),
      ],
    )
    |> simplifile.append(to: out_path)
  Nil
}

fn write_int_to_enum(enum: parser.PbEnum, out_path: String) {
  let head =
    format(
      "
          pub fn int_to_{name}(n: Int) -> {typename} {
            case n {
        ",
      [#("name", pascal_to_snake(enum.name)), #("typename", enum.name)],
    )

  let assert Ok(_) =
    enum.fields
    |> list.map(fn(f) {
      format("{key} -> {value}\n", [
        #("key", int.to_string(f.tag)),
        #("value", f.name),
      ])
    })
    |> list.fold(head, string.append)
    |> string.append(
      "
        _ -> panic
      }}
      ",
    )
    |> simplifile.append(to: out_path)

  Nil
}

fn write_enum_to_int(enum: parser.PbEnum, out_path: String) {
  let head =
    format(
      "
          pub fn {name}_to_int({name}: {typename}) -> Int {
            case {name} {
        ",
      [#("name", pascal_to_snake(enum.name)), #("typename", enum.name)],
    )

  let assert Ok(_) =
    enum.fields
    |> list.map(fn(f) {
      format("{key} -> {value}\n", [
        #("key", f.name),
        #("value", int.to_string(f.tag)),
      ])
    })
    |> list.fold(head, string.append)
    |> string.append(
      "
      }}
      ",
    )
    |> simplifile.append(to: out_path)

  Nil
}

fn get_enums(text: String, lexer, parser) {
  let assert Ok(re) = regex.from_string("enum\\s+\\w+\\s*{[^{}]*}")
  regex.scan(re, text)
  |> list.map(fn(a) {
    let assert Ok(tokens) = lexer.run(a.content, lexer)
    let assert Ok(enum) = nibble.run(tokens, parser)
    enum
  })
}

fn get_structs(text: String, lexer, parser) {
  let assert Ok(re) =
    regex.from_string(
      "//\\s*@gleam\\s+record\\s*\nmessage\\s+\\w+\\s*{([^{}]*)}",
    )
  regex.scan(re, text)
  |> list.map(fn(a) {
    let assert Ok(tokens) = lexer.run(a.content, lexer)
    let assert Ok(message) = nibble.run(tokens, parser)
    #(0, message)
  })
}

fn get_messages(text: String, lexer, parser) {
  let assert Ok(re) =
    regex.from_string(
      "//\\s*@gleam\\s+msgid\\s*=\\s*(\\d+)\\s*\nmessage\\s+\\w+\\s*{([^{}]*)}",
    )
  regex.scan(re, text)
  |> list.map(fn(a) {
    let assert Ok(tokens) = lexer.run(a.content, lexer)
    let assert Ok(msg) = nibble.run(tokens, parser)
    let msgid =
      a.submatches
      |> list.first
      |> result.lazy_unwrap(fn() { panic })
      |> option.lazy_unwrap(fn() { panic })
      |> int.parse
      |> result.lazy_unwrap(fn() { panic })
    #(msgid, msg)
  })
}

// "  Item(id: Int, num: Int)\n"
fn message_to_string(
  message: Message,
  convert: fn(parser.PbMessageField) -> String,
) -> String {
  case list.length(message.fields) > 0 {
    True -> {
      format("  {name}({body})\n", [
        #("name", message.name),
        #("body", {
          message.fields
          |> list.map(convert)
          |> list.fold("", fn(a, b) { a <> b <> ", " })
          |> string.drop_right(2)
        }),
      ])
    }
    False -> "  " <> message.name <> "\n"
  }
}

// varint: int32, int64, uint32, uint64, bool, enum
// i64: fixed64, sfixed64, double
// i32: fixed32, sfixed32, float
fn to_gleam_ty(ty: String, repeated: Bool) -> String {
  let ty = case ty {
    // string
    "string" -> "String"
    // varint
    "int32" | "int64" | "uint32" | "uint64" -> "Int"
    "bool" -> "Bool"
    // i64
    "fixed64" | "sfixed64" | "double" -> "Float"
    // i32
    "fixed32" | "sfixed32" | "float" -> "Float"
    // Custom Type (Struct | Enum)
    x -> x
  }
  case repeated {
    True -> "List(" <> ty <> ")"
    False -> ty
  }
}

fn pascal_to_snake(ident: String) -> String {
  let assert Ok(re) = regex.from_string("[A-Z][a-z]*")
  regex.scan(re, ident)
  |> list.map(fn(a) { a.content })
  |> list.map(fn(a) { string.lowercase(a) })
  |> list.fold("", fn(a, b) { a <> "_" <> b })
  |> string.drop_left(1)
}