package main

import (
	"fmt"
	"strings"

	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/types/pluginpb"
)

func main() {
	protogen.Options{}.Run(func(gen *protogen.Plugin) error {
		gen.SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)
		for _, f := range gen.Files {
			if len(f.Services) == 0 {
				continue
			}
			if err := generateFiles(gen, f); err != nil {
				return err
			}
		}
		return nil
	})
}

func generateFiles(gen *protogen.Plugin, file *protogen.File) error {
	for _, service := range file.Services {
		// Generate server file
		if err := generateServerFile(gen, file, service); err != nil {
			return err
		}

		// Generate client file
		if err := generateClientFile(gen, file, service); err != nil {
			return err
		}

		// Generate logic file
		if err := generateLogicFile(gen, file, service); err != nil {
			return err
		}
	}
	return nil
}

func generateServerFile(gen *protogen.Plugin, file *protogen.File, service *protogen.Service) error {
	filename := fmt.Sprintf("%s_server.pb.go", strings.ToLower(service.GoName))
	g := gen.NewGeneratedFile(filename, file.GoImportPath)

	// Package declaration
	g.P("// Code generated by protoc-gen-slc. DO NOT EDIT.")
	g.P()
	g.P("package ", file.GoPackageName)
	g.P()

	// Imports
	g.P("import (")
	g.P("\t\"context\"")
	g.P("\t\"errors\"")
	g.P()
	g.P("\t\"google.golang.org/grpc\"")
	g.P("\t\"google.golang.org/grpc/codes\"")
	g.P("\t\"google.golang.org/grpc/status\"")
	g.P(")")
	g.P()

	// Server struct
	g.P("type ", service.GoName, "Server struct {")
	g.P("\tUnimplemented", service.GoName, "Server")
	g.P("\tlogic *", service.GoName, "Logic")
	g.P("}")
	g.P()

	// NewServer function
	g.P("func New", service.GoName, "Server(logic *", service.GoName, "Logic) *", service.GoName, "Server {")
	g.P("\treturn &", service.GoName, "Server{logic: logic}")
	g.P("}")
	g.P()

	// Register function
	g.P("func Register", service.GoName, "Server(s *grpc.Server, logic *", service.GoName, "Logic) {")
	g.P("\tserver := New", service.GoName, "Server(logic)")
	g.P("\tRegister", service.GoName, "Server(s, server)")
	g.P("}")
	g.P()

	// Service methods
	for _, method := range service.Methods {
		g.P("func (s *", service.GoName, "Server) ", methodSignature(g, method), " {")
		g.P("\t// Add your server-side logic here")
		g.P("\tresp, err := s.logic.", method.GoName, "(ctx, req)")
		g.P("\tif err != nil {")
		g.P("\t\treturn nil, status.Errorf(codes.Internal, \"%v\", err)")
		g.P("\t}")
		g.P("\treturn resp, nil")
		g.P("}")
		g.P()
	}

	return nil
}

func generateClientFile(gen *protogen.Plugin, file *protogen.File, service *protogen.Service) error {
	filename := fmt.Sprintf("%s_client.pb.go", strings.ToLower(service.GoName))
	g := gen.NewGeneratedFile(filename, file.GoImportPath)

	// Package declaration
	g.P("// Code generated by protoc-gen-layered. DO NOT EDIT.")
	g.P()
	g.P("package ", file.GoPackageName)
	g.P()

	// Imports
	g.P("import (")
	g.P("\t\"context\"")
	g.P()
	g.P("\t\"google.golang.org/grpc\"")
	g.P(")")
	g.P()

	// Client struct
	g.P("type ", service.GoName, "Client struct {")
	g.P("\tcc grpc.ClientConnInterface")
	g.P("}")
	g.P()

	// NewClient function
	g.P("func New", service.GoName, "Client(cc grpc.ClientConnInterface) *", service.GoName, "Client {")
	g.P("\treturn &", service.GoName, "Client{cc}")
	g.P("}")
	g.P()

	// Client methods
	for _, method := range service.Methods {
		g.P("func (c *", service.GoName, "Client) ", methodSignature(g, method), " {")
		g.P("\tout := new(", method.Output.GoIdent, ")")
		g.P("\terr := c.cc.Invoke(ctx, \"", fullMethodName(file, service, method), "\", req, out)")
		g.P("\tif err != nil {")
		g.P("\t\treturn nil, err")
		g.P("\t}")
		g.P("\treturn out, nil")
		g.P("}")
		g.P()
	}

	return nil
}

func generateLogicFile(gen *protogen.Plugin, file *protogen.File, service *protogen.Service) error {
	filename := fmt.Sprintf("%s_logic.pb.go", strings.ToLower(service.GoName))
	g := gen.NewGeneratedFile(filename, file.GoImportPath)

	// Package declaration
	g.P("// Code generated by protoc-gen-layered. DO NOT EDIT.")
	g.P()
	g.P("package ", file.GoPackageName)
	g.P()

	// Imports
	g.P("import (")
	g.P("\t\"context\"")
	g.P("\t\"errors\"")
	g.P(")")
	g.P()

	// Logic struct
	g.P("type ", service.GoName, "Logic struct {")
	g.P("\t// Add your dependencies here")
	g.P("}")
	g.P()

	// NewLogic function
	g.P("func New", service.GoName, "Logic() *", service.GoName, "Logic {")
	g.P("\treturn &", service.GoName, "Logic{}")
	g.P("}")
	g.P()

	// Logic methods
	for _, method := range service.Methods {
		g.P("func (l *", service.GoName, "Logic) ", method.GoName, "(ctx context.Context, req *", method.Input.GoIdent, ") (*", method.Output.GoIdent, ", error) {")
		g.P("\t// Implement your business logic here")
		g.P("\treturn nil, errors.New(\"not implemented\")")
		g.P("}")
		g.P()
	}

	return nil
}

func methodSignature(g *protogen.GeneratedFile, method *protogen.Method) string {
	return fmt.Sprintf("%s(ctx context.Context, req *%s) (*%s, error)",
		method.GoName,
		method.Input.GoIdent,
		method.Output.GoIdent)
}

func fullMethodName(file *protogen.File, service *protogen.Service, method *protogen.Method) string {
	return fmt.Sprintf("/%s.%s/%s",
		file.Proto.GetPackage(),
		service.GoName,
		method.GoName)
}