// Copyright 2018 Twitch Interactive, Inc.  All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may not
// use this file except in compliance with the License. A copy of the License is
// located at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// or in the "license" file accompanying this file. This file is distributed on
// an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
// express or implied. See the License for the specific language governing
// permissions and limitations under the License.

package main

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"unicode"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	plugin "github.com/golang/protobuf/protoc-gen-go/plugin"
)

func main() {
	genReq := readGenRequest(os.Stdin)
	g := &generator{version: Version, genReq: genReq}
	genResp := g.Generate()
	writeGenResponse(os.Stdout, genResp)
}

type generator struct {
	version string
	genReq  *plugin.CodeGeneratorRequest
}

func (g *generator) Generate() *plugin.CodeGeneratorResponse {
	resp := new(plugin.CodeGeneratorResponse)

	for _, f := range g.protoFilesToGenerate() {
		twirpFileName := noExtension(filePath(f)) + "_twirp.rb"             // e.g. "hello_world/service_twirp.rb"
		pbFileRelativePath := noExtension(onlyBase(filePath(f))) + "_pb.rb" // e.g. "service_pb.rb"

		rubyCode := g.generateRubyCode(f, pbFileRelativePath)
		respFile := &plugin.CodeGeneratorResponse_File{
			Name:    proto.String(twirpFileName),
			Content: proto.String(rubyCode),
		}
		resp.File = append(resp.File, respFile)
	}

	return resp
}

func (g *generator) generateRubyCode(file *descriptor.FileDescriptorProto, pbFileRelativePath string) string {
	b := new(bytes.Buffer)
	print(b, "# Code generated by protoc-gen-twirp_ruby %s, DO NOT EDIT.", g.version)
	print(b, "require 'twirp'")
	print(b, "require_relative '%s'", pbFileRelativePath) // require generated file with messages
	print(b, "")

	indent := indentation(0)
	pkgName := file.GetPackage()
	modules := packageToRubyModules(pkgName)
	for _, m := range modules {
		print(b, "%smodule %s", indent, m)
		indent += 1
	}

	for i, service := range file.Service {
		svcName := service.GetName()

		print(b, "%sclass %sService < Twirp::Service", indent, camelCase(svcName))
		if pkgName != "" {
			print(b, "%s  package '%s'", indent, pkgName)
		}
		print(b, "%s  service '%s'", indent, svcName)
		for _, method := range service.GetMethod() {
			rpcName := method.GetName()
			print(b, "%s  rpc :%s, %s, %s, :ruby_method => :%s",
				indent, rpcName, methodInputName(method), methodOutputName(method), snakeCase(rpcName))
		}
		print(b, "%send", indent)
		print(b, "")

		print(b, "%sclass %sClient < Twirp::Client", indent, camelCase(svcName))
		print(b, "%s  client_for %sService", indent, camelCase(svcName))
		print(b, "%send", indent)
		if i < len(file.Service)-1 {
			print(b, "")
		}
	}

	for range modules {
		indent -= 1
		print(b, "%send", indent)
	}

	return b.String()
}

// protoFilesToGenerate selects descriptor proto files that were explicitly listed on the command-line.
func (g *generator) protoFilesToGenerate() []*descriptor.FileDescriptorProto {
	files := []*descriptor.FileDescriptorProto{}
	for _, name := range g.genReq.FileToGenerate { // explicitly listed on the command-line
		for _, f := range g.genReq.ProtoFile { // all files and everything they import
			if f.GetName() == name { // match
				files = append(files, f)
				continue
			}
		}
	}
	return files
}

// indentation represents the level of Ruby indentation for a block of code. It
// implements the fmt.Stringer interface to output the correct number of spaces
// for the given level of indentation
type indentation int

func (i indentation) String() string {
	return strings.Repeat("  ", int(i))
}

func print(buf *bytes.Buffer, tpl string, args ...interface{}) {
	buf.WriteString(fmt.Sprintf(tpl, args...))
	buf.WriteByte('\n')
}

func filePath(f *descriptor.FileDescriptorProto) string {
	return *f.Name
}

func onlyBase(path string) string {
	return filepath.Base(path)
}

func noExtension(path string) string {
	ext := filepath.Ext(path)
	return strings.TrimSuffix(path, ext)
}

func methodInputName(meth *descriptor.MethodDescriptorProto) string {
	fullName := meth.GetInputType()
	split := strings.Split(fullName, ".")
	return split[len(split)-1]
}

func methodOutputName(meth *descriptor.MethodDescriptorProto) string {
	fullName := meth.GetOutputType()
	split := strings.Split(fullName, ".")
	return split[len(split)-1]
}

func Fail(msgs ...string) {
	s := strings.Join(msgs, " ")
	log.Print("error:", s)
	os.Exit(1)
}

func readGenRequest(r io.Reader) *plugin.CodeGeneratorRequest {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		Fail(err.Error(), "reading input")
	}

	req := new(plugin.CodeGeneratorRequest)
	if err = proto.Unmarshal(data, req); err != nil {
		Fail(err.Error(), "parsing input proto")
	}

	if len(req.FileToGenerate) == 0 {
		Fail("no files to generate")
	}

	return req
}

func writeGenResponse(w io.Writer, resp *plugin.CodeGeneratorResponse) {
	data, err := proto.Marshal(resp)
	if err != nil {
		Fail(err.Error(), "marshaling response")
	}
	_, err = w.Write(data)
	if err != nil {
		Fail(err.Error(), "writing response")
	}
}

// Modules converts protobuf package name to a list of Ruby module names to
// represent it. e.g. packageToRubyModules("my.cool.package") => ["My", "Cool", "Package"]
func packageToRubyModules(pkgName string) []string {
	parts := []string{}
	for _, p := range strings.Split(pkgName, ".") {
		parts = append(parts, camelCase(p))
	}
	return parts
}

// snakeCase converts a string from CamelCase to snake_case.
func snakeCase(s string) string {
	var buf bytes.Buffer
	for i, r := range s {
		if unicode.IsUpper(r) && i > 0 {
			fmt.Fprintf(&buf, "_")
		}
		r = unicode.ToLower(r)
		fmt.Fprintf(&buf, "%c", r)
	}
	return buf.String()
}

// camelCase converts a string from snake_case to CamelCased.
// If there is an interior underscore followed by a lower case letter, drop the
// underscore and convert the letter to upper case. There is a remote
// possibility of this rewrite causing a name collision, but it's so remote
// we're prepared to pretend it's nonexistent - since the C++ generator
// lowercases names, it's extremely unlikely to have two fields with different
// capitalizations. In short, _my_field_name_2 becomes XMyFieldName_2.
func camelCase(s string) string {
	if s == "" {
		return ""
	}
	t := make([]byte, 0, 32)
	i := 0
	if s[0] == '_' {
		// Need a capital letter; drop the '_'.
		t = append(t, 'X')
		i++
	}
	// Invariant: if the next letter is lower case, it must be converted
	// to upper case.
	//
	// That is, we process a word at a time, where words are marked by _ or upper
	// case letter. Digits are treated as words.
	for ; i < len(s); i++ {
		c := s[i]
		if c == '_' && i+1 < len(s) && (isASCIILower(s[i+1]) || isASCIIDigit(s[i+1])) {
			continue // Skip the underscore in s.
		}
		// Assume we have a letter now - if not, it's a bogus identifier. The next
		// word is a sequence of characters that must start upper case.
		if isASCIILower(c) {
			c ^= ' ' // Make it a capital letter.
		}
		t = append(t, c) // Guaranteed not lower case.
		// Accept lower case sequence that follows.
		for i+1 < len(s) && (isASCIILower(s[i+1]) || isASCIIDigit(s[i+1])) {
			i++
			t = append(t, s[i])
		}
	}
	return string(t)
}

// Is c an ASCII lower-case letter?
func isASCIILower(c byte) bool {
	return 'a' <= c && c <= 'z'
}

// Is c an ASCII digit?
func isASCIIDigit(c byte) bool {
	return '0' <= c && c <= '9'
}
