// The gRPC service tree, goctl's `rpc new` and `rpc protoc` in MoonBit: the
// generated messages and handler tables under `pb/`, the author's logic under
// `internal/logic`, the wiring between them under `internal/server`, and a typed
// client under `client/`.

///|
/// `t` as the type of a request or reply seen from outside `pb`.
fn pb(t : String, many : Bool) -> String {
  if many {
    "Array[@pb." + t + "]"
  } else {
    "@pb." + t
  }
}

///|
/// The module a proto's tree is written as: its package, or the first service's
/// name when it declares none.
fn proto_module(proto : @proto.Proto) -> String {
  if proto.package_ != "" {
    proto.package_.replace_all(old=".", new="_")
  } else if proto.services.length() > 0 {
    @text.to_snake(proto.services[0].name)
  } else {
    "service"
  }
}

///|
/// The logic seeded for an RPC: a function of the context and the request that
/// says it is not implemented yet, as gRPC's generated servers do.
fn logic_src(st : @style.Style, r : @proto.Rpc, body : String) -> String {
  let src =
    #|///|
    #|/// `%rpc`. This is where the service does its work.
    #|pub fn %fn(
    #|  _svc : @svc.ServiceContext,
    #|  _req : %in,
    #|) -> Result[%out, @moonrpc.Status] {
    #|  %body
    #|}
  src
  .replace_all(old="%rpc", new=r.name)
  .replace_all(old="%fn", new=st.format(r.name + "_logic"))
  .replace_all(old="%in", new=pb(r.request, r.client_streaming))
  .replace_all(old="%out", new=pb(r.response, r.server_streaming))
  .replace_all(old="%body", new=body) +
  "\n"
}

///|
/// `internal/server`: each service's record filled with the logic functions.
fn server_src(st : @style.Style, proto : @proto.Proto) -> String {
  let out = StringBuilder()
  out.write_string("// Code generated by moonctl. DO NOT EDIT.\n")
  for s in proto.services {
    out.write_string(
      "\n///|\n/// `" +
      s.name +
      "` with each RPC answered by its function in `internal/logic`.\npub fn " +
      st.format(s.name + "_server") +
      "(svc : @svc.ServiceContext) -> @pb." +
      s.name +
      "Server {\n  {\n",
    )
    for r in s.rpcs {
      out.write_string(
        "    " +
        @text.to_snake(r.name) +
        ": req => @logic." +
        st.format(r.name + "_logic") +
        "(svc, req),\n",
      )
    }
    out.write_string("  }\n}\n")
  }
  out.to_string()
}

///|
/// `client/`: per service a struct over a channel with one typed call per RPC. A
/// non-OK status raises `CallError` with the status code and message.
fn client_src(proto : @proto.Proto) -> String {
  let head =
    #|// Code generated by moonctl. DO NOT EDIT.
    #|
    #|///|
    #|/// A call the server ended with a non-OK gRPC status.
    #|pub(all) suberror CallError {
    #|  CallError(code~ : Int32, message~ : String)
    #|} derive(Debug)
    #|
    #|///|
    #|pub extend CallError with Debug::{to_repr}
    #|
    #|///|
    #|fn check(reply : @moonrpc.CallReply) -> Unit raise CallError {
    #|  if reply.grpc_status != 0 {
    #|    raise CallError(code=reply.grpc_status, message=reply.grpc_message)
    #|  }
    #|}
    #|
  let out = StringBuilder()
  out.write_string(head)
  for s in proto.services {
    let full = if proto.package_ == "" {
      s.name
    } else {
      proto.package_ + "." + s.name
    }
    let client = s.name + "Client"
    out.write_string(
      "\n///|\n/// A client of `" +
      full +
      "`.\npub struct " +
      client +
      " {\n  channel : @net.Channel\n}\n\n///|\npub fn " +
      client +
      "::new(channel : @net.Channel) -> " +
      client +
      " {\n  { channel, }\n}\n\n///|\npub fn " +
      client +
      "::close(self : " +
      client +
      ") -> Unit {\n  self.channel.close()\n}\n",
    )
    for r in s.rpcs {
      let path = @text.quote("/" + full + "/" + r.name)
      let (call, arg) = match (r.client_streaming, r.server_streaming) {
        (false, false) => ("unary", "req.encode()")
        (false, true) => ("server_streaming", "req.encode()")
        (true, false) => ("client_streaming", "reqs.map(req => req.encode())")
        (true, true) => ("bidi", "reqs.map(req => req.encode())")
      }
      let param = if r.client_streaming { "reqs" } else { "req" }
      let result = if r.server_streaming {
        "reply.messages.map(@pb." + r.response + "::decode)"
      } else {
        "@pb." + r.response + "::decode(reply.messages.get(0).unwrap_or(b\"\"))"
      }
      out.write_string(
        "\n///|\n/// `" +
        r.name +
        "`.\npub async fn " +
        client +
        "::" +
        @text.to_snake(r.name) +
        "(\n  self : " +
        client +
        ",\n  " +
        param +
        " : " +
        pb(r.request, r.client_streaming) +
        ",\n) -> " +
        pb(r.response, r.server_streaming) +
        " raise {\n  let reply = self.channel." +
        call +
        "(" +
        path +
        ", " +
        arg +
        ")\n  check(reply)\n  " +
        result +
        "\n}\n",
      )
    }
  }
  out.to_string()
}

///|
/// Registering every service of `proto` on `grpc`, as the entry point and the
/// client test both do.
fn register_src(
  st : @style.Style,
  proto : @proto.Proto,
  indent : String,
) -> String {
  let out = StringBuilder()
  for s in proto.services {
    out.write_string(
      indent +
      "for path, handler in @server." +
      st.format(s.name + "_server") +
      "(svc).handlers() {\n" +
      indent +
      "  grpc.handlers[path] = handler\n" +
      indent +
      "}\n",
    )
  }
  out.to_string()
}

///|
/// The gRPC tree for `proto`. `logic` seeds the body of an RPC's logic by name;
/// one not in it starts as `Err(Unimplemented)`. `smoke` is the client test to
/// seed, when there is one.
fn rpc_tree(
  proto : @proto.Proto,
  st : @style.Style,
  dir : String,
  logic? : Map[String, String],
  smoke? : String,
) -> Array[TreeFile] {
  let logic = logic.unwrap_or(Map([]))
  let mod_ = proto_module(proto)
  let file = st.format(mod_)
  let pb_pkg = mod_ + "/pb"
  let svc_pkg = mod_ + "/internal/svc"
  let config_pkg = mod_ + "/internal/config"
  let logic_pkg = mod_ + "/internal/logic"
  let server_pkg = mod_ + "/internal/server"
  let rpc = ("moonbitstack/moonrpc", "")
  let net = ("moonbitstack/moonrpc/net", "")
  let main_src =
    #|///|
    #|/// The %svc service: read the configuration, build the context its logic
    #|/// shares and serve every RPC over gRPC.
    #|async fn main {
    #|  let cfg = @config.Config::new()
    #|  let svc = @svc.ServiceContext::new(cfg)
    #|  let grpc = @net.GrpcServer::new()
    #|%register  println("%svc serving gRPC on " + cfg.host + ":" + cfg.port.to_string())
    #|  grpc.serve(host=cfg.host, port=cfg.port)
    #|}
  let out = base_tree(
    mod_,
    st,
    dir,
    deps=["moonbitstack/moonrpc"],
    imports=[net, (server_pkg, "")],
    main=main_src
      .replace_all(old="%svc", new=mod_)
      .replace_all(old="%register", new=register_src(st, proto, "  ")) +
      "\n",
  )
  out.push({
    path: at(dir, "pb/moon.pkg"),
    content: pkg_file([rpc]),
    regen: Always,
  })
  out.push({
    path: at(dir, "pb/" + file + "_grpc.mbt"),
    content: @proto.generate_grpc(proto),
    regen: Always,
  })
  out.push({
    path: at(dir, "internal/logic/moon.pkg"),
    content: pkg_file([rpc, (svc_pkg, ""), (pb_pkg, "")]),
    regen: Once,
  })
  for s in proto.services {
    for r in s.rpcs {
      out.push({
        path: at(dir, "internal/logic/" + st.format(r.name + "_logic") + ".mbt"),
        content: logic_src(
          st,
          r,
          logic.get(r.name).unwrap_or("Err(Unimplemented)"),
        ),
        regen: Once,
      })
    }
  }
  out.push({
    path: at(dir, "internal/server/moon.pkg"),
    content: pkg_file([(pb_pkg, ""), (svc_pkg, ""), (logic_pkg, "")]),
    regen: Always,
  })
  out.push({
    path: at(dir, "internal/server/" + st.format(mod_ + "_server") + ".mbt"),
    content: server_src(st, proto),
    regen: Always,
  })
  let client_test = pkg_file([
    ("moonbitlang/async", ""),
    (config_pkg, ""),
    (svc_pkg, ""),
    (server_pkg, ""),
  ]).replace(old="}\n", new="} for \"test\"\n")
  out.push({
    path: at(dir, "client/moon.pkg"),
    content: "supported_targets = \"native\"\n\n" +
    pkg_file([rpc, net, (pb_pkg, "")]) +
    (if smoke is Some(_) { "\n" + client_test } else { "" }),
    regen: Once,
  })
  out.push({
    path: at(dir, "client/" + st.format(mod_ + "_client") + ".mbt"),
    content: client_src(proto),
    regen: Always,
  })
  if smoke is Some(src) {
    out.push({
      path: at(dir, "client/" + st.format(mod_ + "_client") + "_test.mbt"),
      content: src.replace_all(
        old="%register",
        new=register_src(st, proto, "    "),
      ),
      regen: Once,
    })
  }
  out
}

///|
/// The layered gRPC tree goctl writes for a `.proto`, in MoonBit, under `dir`:
/// `pb/` with the messages, codecs and handler tables (rewritten on every run),
/// `internal/logic` with one function per RPC and `internal/config`/`internal/svc`
/// (seeded once), `internal/server` wiring the logic into each service, `client/`
/// with a typed client per service, and the entry point serving it all on 12000.
///
/// Hand the result to `tree_plan` to keep what the author has filled in.
pub fn generate_rpc_tree(
  proto : @proto.Proto,
  style? : @style.Style,
  dir~ : String = "",
) -> Array[TreeFile] {
  rpc_tree(proto, style.unwrap_or(@style.Style::gozero()), dir)
}