// Copyright AGNTCY Contributors (https://github.com/agntcy)
// SPDX-License-Identifier: Apache-2.0

use anyhow::{Context, Result};
use prost_types::compiler::{
    CodeGeneratorRequest, CodeGeneratorResponse, code_generator_response::File,
};
use std::collections::HashMap;

// --- TEMPLATE DEFINITIONS ---

const FILE_TEMPLATE: &str = r#"// Code generated by protoc-gen-slimrpc-go. DO NOT EDIT.
// source: {{PROTO_FILE}}

package {{PACKAGE}}

import (
	"context"
	"fmt"
	"time"

	slim_bindings "github.com/agntcy/slim-bindings-go"
	"github.com/agntcy/slim-bindings-go/slimrpc"
	"google.golang.org/protobuf/proto"
{{PROTO_IMPORTS}}
)

{{SERVICE_DEFINITIONS}}
"#;

const CLIENT_INTERFACE_TEMPLATE: &str = r#"
// {{SERVICE_NAME}}Client is the client API for {{SERVICE_NAME}} service.
type {{SERVICE_NAME}}Client interface {
{{CLIENT_METHODS}}
}

type {{SERVICE_NAME}}ClientImpl struct {
	channel *slim_bindings.Channel
}

// New{{SERVICE_NAME}}Client creates a new {{SERVICE_NAME}} client.
func New{{SERVICE_NAME}}Client(channel *slim_bindings.Channel) {{SERVICE_NAME}}Client {
	return &{{SERVICE_NAME}}ClientImpl{
		channel: channel,
	}
}

{{CLIENT_METHOD_IMPLS}}
"#;

const UNARY_UNARY_CLIENT_METHOD: &str = r#"	{{METHOD_NAME}}(ctx context.Context, req *{{INPUT_TYPE}}) (*{{OUTPUT_TYPE}}, error)
"#;

const UNARY_UNARY_CLIENT_IMPL: &str = r#"
func (c *{{SERVICE_NAME}}ClientImpl) {{METHOD_NAME}}(ctx context.Context, req *{{INPUT_TYPE}}) (*{{OUTPUT_TYPE}}, error) {
	// Serialize request
	reqBytes, err := proto.Marshal(req)
	if err != nil {
		return nil, err
	}

	// Extract timeout from context
	var timeout *time.Duration
	if deadline, ok := ctx.Deadline(); ok {
		t := time.Until(deadline)
		timeout = &t
	}

	// Extract metadata from context
	var metadata *map[string]string
	if md, ok := slimrpc.MetadataFromContext(ctx); ok {
		metadata = &md
	}

	// Make RPC call
	respBytes, err := c.channel.CallUnaryAsync("{{PACKAGE_NAME}}.{{SERVICE_NAME}}", "{{METHOD_NAME}}", reqBytes, timeout, metadata)
	if err != nil {
		return nil, err
	}

	// Deserialize response
	resp := &{{OUTPUT_TYPE}}{}
	if err := proto.Unmarshal(respBytes, resp); err != nil {
		return nil, err
	}

	return resp, nil
}
"#;

const UNARY_STREAM_CLIENT_METHOD: &str = r#"	{{METHOD_NAME}}(ctx context.Context, req *{{INPUT_TYPE}}) (slimrpc.ResponseStream[*{{OUTPUT_TYPE}}], error)
"#;

const UNARY_STREAM_CLIENT_IMPL: &str = r#"
func (c *{{SERVICE_NAME}}ClientImpl) {{METHOD_NAME}}(ctx context.Context, req *{{INPUT_TYPE}}) (slimrpc.ResponseStream[*{{OUTPUT_TYPE}}], error) {
	// Serialize request
	reqBytes, err := proto.Marshal(req)
	if err != nil {
		return nil, err
	}

	// Extract timeout from context
	var timeout *time.Duration
	if deadline, ok := ctx.Deadline(); ok {
		t := time.Until(deadline)
		timeout = &t
	}

	// Extract metadata from context
	var metadata *map[string]string
	if md, ok := slimrpc.MetadataFromContext(ctx); ok {
		metadata = &md
	}

	// Make RPC call
	stream, err := c.channel.CallUnaryStreamAsync("{{PACKAGE_NAME}}.{{SERVICE_NAME}}", "{{METHOD_NAME}}", reqBytes, timeout, metadata)
	if err != nil {
		return nil, err
	}

	return slimrpc.NewClientResponseStream[*{{OUTPUT_TYPE}}](stream), nil
}
"#;

const STREAM_UNARY_CLIENT_METHOD: &str = r#"	{{METHOD_NAME}}(ctx context.Context) (slimrpc.ClientRequestStream[*{{INPUT_TYPE}}, *{{OUTPUT_TYPE}}], error)
"#;

const STREAM_UNARY_CLIENT_IMPL: &str = r#"
func (c *{{SERVICE_NAME}}ClientImpl) {{METHOD_NAME}}(ctx context.Context) (slimrpc.ClientRequestStream[*{{INPUT_TYPE}}, *{{OUTPUT_TYPE}}], error) {
	// Extract timeout from context
	var timeout *time.Duration
	if deadline, ok := ctx.Deadline(); ok {
		t := time.Until(deadline)
		timeout = &t
	}

	// Extract metadata from context
	var metadata *map[string]string
	if md, ok := slimrpc.MetadataFromContext(ctx); ok {
		metadata = &md
	}

	stream := c.channel.CallStreamUnary("{{PACKAGE_NAME}}.{{SERVICE_NAME}}", "{{METHOD_NAME}}", timeout, metadata)
	return slimrpc.NewClientRequestStream[*{{INPUT_TYPE}}, *{{OUTPUT_TYPE}}](stream), nil
}
"#;

const STREAM_STREAM_CLIENT_METHOD: &str = r#"	{{METHOD_NAME}}(ctx context.Context) (slimrpc.ClientBidiStream[*{{INPUT_TYPE}}, *{{OUTPUT_TYPE}}], error)
"#;

const STREAM_STREAM_CLIENT_IMPL: &str = r#"
func (c *{{SERVICE_NAME}}ClientImpl) {{METHOD_NAME}}(ctx context.Context) (slimrpc.ClientBidiStream[*{{INPUT_TYPE}}, *{{OUTPUT_TYPE}}], error) {
	// Extract timeout from context
	var timeout *time.Duration
	if deadline, ok := ctx.Deadline(); ok {
		t := time.Until(deadline)
		timeout = &t
	}

	// Extract metadata from context
	var metadata *map[string]string
	if md, ok := slimrpc.MetadataFromContext(ctx); ok {
		metadata = &md
	}

	stream := c.channel.CallStreamStream("{{PACKAGE_NAME}}.{{SERVICE_NAME}}", "{{METHOD_NAME}}", timeout, metadata)
	return slimrpc.NewClientBidiStream[*{{INPUT_TYPE}}, *{{OUTPUT_TYPE}}](stream), nil
}
"#;

const SERVER_INTERFACE_TEMPLATE: &str = r#"
// {{SERVICE_NAME}}Server is the server API for {{SERVICE_NAME}} service.
// All implementations must embed Unimplemented{{SERVICE_NAME}}Server
// for forward compatibility
type {{SERVICE_NAME}}Server interface {
{{SERVER_METHODS}}
}

// Unimplemented{{SERVICE_NAME}}Server must be embedded to have forward compatible implementations.
type Unimplemented{{SERVICE_NAME}}Server struct {
}

{{UNIMPLEMENTED_METHODS}}

// Register{{SERVICE_NAME}}Server registers the server with slim_bindings.
func Register{{SERVICE_NAME}}Server(server *slim_bindings.Server, impl {{SERVICE_NAME}}Server) {
{{REGISTER_METHODS}}
}

{{HANDLER_IMPLS}}
"#;

const UNIMPLEMENTED_METHOD: &str = r#"
func (Unimplemented{{SERVICE_NAME}}Server) {{METHOD_NAME}}(ctx context.Context, req *{{INPUT_TYPE}}) (*{{OUTPUT_TYPE}}, error) {
	return nil, fmt.Errorf("method {{METHOD_NAME}} not implemented")
}
"#;

const UNARY_UNARY_HANDLER_IMPL: &str = r#"
type {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler struct {
	impl {{SERVICE_NAME}}Server
}

func (h *{{SERVICE_NAME}}_{{METHOD_NAME}}_Handler) Handle(request []byte, rpcContext *slim_bindings.Context) ([]byte, error) {
	req := &{{INPUT_TYPE}}{}
	if err := proto.Unmarshal(request, req); err != nil {
		return nil, slim_bindings.NewRpcErrorRpc(
			slim_bindings.RpcCodeInvalidArgument,
			err.Error(),
			nil,
		)
	}

	// Convert slim_bindings.Context to context.Context
	ctx, cancel := slimrpc.ContextFromRpcContext(rpcContext)
	defer cancel()

	resp, err := h.impl.{{METHOD_NAME}}(ctx, req)
	if err != nil {
		// Check if it's already an RpcError
		if rpcErr, ok := err.(*slim_bindings.RpcError); ok {
			return nil, rpcErr
		}
		// Convert generic errors to RpcError
		return nil, slim_bindings.NewRpcErrorRpc(
			slim_bindings.RpcCodeInternal,
			err.Error(),
			nil,
		)
	}

	respBytes, err := proto.Marshal(resp)
	if err != nil {
		return nil, slim_bindings.NewRpcErrorRpc(
			slim_bindings.RpcCodeInternal,
			err.Error(),
			nil,
		)
	}
	return respBytes, nil
}
"#;

const REGISTER_METHOD: &str = r#"	server.RegisterUnaryUnary("{{PACKAGE_NAME}}.{{SERVICE_NAME}}", "{{METHOD_NAME}}", &{{SERVICE_NAME}}_{{METHOD_NAME}}_Handler{impl: impl})
"#;

const UNARY_UNARY_SERVER_METHOD: &str = r#"	{{METHOD_NAME}}(ctx context.Context, req *{{INPUT_TYPE}}) (*{{OUTPUT_TYPE}}, error)
"#;

const UNARY_STREAM_SERVER_METHOD: &str = r#"	{{METHOD_NAME}}(ctx context.Context, req *{{INPUT_TYPE}}, stream slimrpc.RequestStream[*{{OUTPUT_TYPE}}]) error
"#;

const STREAM_UNARY_SERVER_METHOD: &str = r#"	{{METHOD_NAME}}(ctx context.Context, stream slimrpc.ResponseStream[*{{INPUT_TYPE}}]) (*{{OUTPUT_TYPE}}, error)
"#;

const STREAM_STREAM_SERVER_METHOD: &str = r#"	{{METHOD_NAME}}(ctx context.Context, stream slimrpc.ServerBidiStream[*{{INPUT_TYPE}}, *{{OUTPUT_TYPE}}]) error
"#;

const UNIMPLEMENTED_UNARY_STREAM_METHOD: &str = r#"
func (Unimplemented{{SERVICE_NAME}}Server) {{METHOD_NAME}}(ctx context.Context, req *{{INPUT_TYPE}}, stream slimrpc.RequestStream[*{{OUTPUT_TYPE}}]) error {
	return fmt.Errorf("method {{METHOD_NAME}} not implemented")
}
"#;

const UNIMPLEMENTED_STREAM_UNARY_METHOD: &str = r#"
func (Unimplemented{{SERVICE_NAME}}Server) {{METHOD_NAME}}(ctx context.Context, stream slimrpc.ResponseStream[*{{INPUT_TYPE}}]) (*{{OUTPUT_TYPE}}, error) {
	return nil, fmt.Errorf("method {{METHOD_NAME}} not implemented")
}
"#;

const UNIMPLEMENTED_STREAM_STREAM_METHOD: &str = r#"
func (Unimplemented{{SERVICE_NAME}}Server) {{METHOD_NAME}}(ctx context.Context, stream slimrpc.ServerBidiStream[*{{INPUT_TYPE}}, *{{OUTPUT_TYPE}}]) error {
	return fmt.Errorf("method {{METHOD_NAME}} not implemented")
}
"#;

const UNARY_STREAM_HANDLER_IMPL: &str = r#"
type {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler struct {
	impl {{SERVICE_NAME}}Server
}

func (h *{{SERVICE_NAME}}_{{METHOD_NAME}}_Handler) Handle(request []byte, rpcContext *slim_bindings.Context, sink *slim_bindings.ResponseSink) error {
	req := &{{INPUT_TYPE}}{}
	if err := proto.Unmarshal(request, req); err != nil {
		rpcErr := slim_bindings.NewRpcErrorRpc(
			slim_bindings.RpcCodeInvalidArgument,
			err.Error(),
			nil,
		)
		sink.SendErrorAsync(rpcErr)
		return rpcErr
	}

	// Convert slim_bindings.Context to context.Context
	ctx, cancel := slimrpc.ContextFromRpcContext(rpcContext)
	defer cancel()

	stream := slimrpc.NewServerRequestStream[*{{OUTPUT_TYPE}}](sink)
	err := h.impl.{{METHOD_NAME}}(ctx, req, stream)

	// Close the stream after handler returns
	closeErr := sink.CloseAsync()
	if err == nil {
		err = closeErr
	}

	if err != nil {
		// Check if it's already an RpcError
		if rpcErr, ok := err.(*slim_bindings.RpcError); ok {
			sink.SendErrorAsync(rpcErr)
			return rpcErr
		}
		// Convert generic errors to RpcError
		rpcErr := slim_bindings.NewRpcErrorRpc(
			slim_bindings.RpcCodeInternal,
			err.Error(),
			nil,
		)
		sink.SendErrorAsync(rpcErr)
		return rpcErr
	}

	return nil
}
"#;

const STREAM_UNARY_HANDLER_IMPL: &str = r#"
type {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler struct {
	impl {{SERVICE_NAME}}Server
}

func (h *{{SERVICE_NAME}}_{{METHOD_NAME}}_Handler) Handle(stream *slim_bindings.RequestStream, rpcContext *slim_bindings.Context) ([]byte, error) {
	// Convert slim_bindings.Context to context.Context
	ctx, cancel := slimrpc.ContextFromRpcContext(rpcContext)
	defer cancel()

	responseStream := slimrpc.NewServerResponseStream[*{{INPUT_TYPE}}](stream)

	resp, err := h.impl.{{METHOD_NAME}}(ctx, responseStream)
	if err != nil {
		if rpcErr, ok := err.(*slim_bindings.RpcError); ok {
			return nil, rpcErr
		}
		return nil, slim_bindings.NewRpcErrorRpc(
			slim_bindings.RpcCodeInternal,
			err.Error(),
			nil,
		)
	}

	respBytes, err := proto.Marshal(resp)
	if err != nil {
		return nil, slim_bindings.NewRpcErrorRpc(
			slim_bindings.RpcCodeInternal,
			err.Error(),
			nil,
		)
	}
	return respBytes, nil
}
"#;

const STREAM_STREAM_HANDLER_IMPL: &str = r#"
type {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler struct {
	impl {{SERVICE_NAME}}Server
}

func (h *{{SERVICE_NAME}}_{{METHOD_NAME}}_Handler) Handle(stream *slim_bindings.RequestStream, rpcContext *slim_bindings.Context, sink *slim_bindings.ResponseSink) error {
	// Convert slim_bindings.Context to context.Context
	ctx, cancel := slimrpc.ContextFromRpcContext(rpcContext)
	defer cancel()

	serverStream := slimrpc.NewServerBidiStream[*{{INPUT_TYPE}}, *{{OUTPUT_TYPE}}](stream, sink)
	err := h.impl.{{METHOD_NAME}}(ctx, serverStream)

	// Close the stream after handler returns
	closeErr := sink.CloseAsync()
	if err == nil {
		err = closeErr
	}

	if err != nil {
		// Check if it's already an RpcError
		if rpcErr, ok := err.(*slim_bindings.RpcError); ok {
			sink.SendErrorAsync(rpcErr)
			return rpcErr
		}
		// Convert generic errors to RpcError
		rpcErr := slim_bindings.NewRpcErrorRpc(
			slim_bindings.RpcCodeInternal,
			err.Error(),
			nil,
		)
		sink.SendErrorAsync(rpcErr)
		return rpcErr
	}

	return nil
}
"#;

const REGISTER_UNARY_STREAM_METHOD: &str = r#"	server.RegisterUnaryStream("{{PACKAGE_NAME}}.{{SERVICE_NAME}}", "{{METHOD_NAME}}", &{{SERVICE_NAME}}_{{METHOD_NAME}}_Handler{impl: impl})
"#;

const REGISTER_STREAM_UNARY_METHOD: &str = r#"	server.RegisterStreamUnary("{{PACKAGE_NAME}}.{{SERVICE_NAME}}", "{{METHOD_NAME}}", &{{SERVICE_NAME}}_{{METHOD_NAME}}_Handler{impl: impl})
"#;

const REGISTER_STREAM_STREAM_METHOD: &str = r#"	server.RegisterStreamStream("{{PACKAGE_NAME}}.{{SERVICE_NAME}}", "{{METHOD_NAME}}", &{{SERVICE_NAME}}_{{METHOD_NAME}}_Handler{impl: impl})
"#;

// --- END TEMPLATE DEFINITIONS ---

/// Generate Go slimrpc code from a CodeGeneratorRequest
pub fn generate(request: CodeGeneratorRequest) -> Result<CodeGeneratorResponse> {
    let mut response = CodeGeneratorResponse {
        supported_features: Some(1), // FEATURE_PROTO3_OPTIONAL
        ..Default::default()
    };

    // Build a map from file name to package
    let mut file_to_package: HashMap<String, String> = HashMap::new();
    for file_desc in &request.proto_file {
        if let (Some(name), Some(package)) = (&file_desc.name, &file_desc.package) {
            file_to_package.insert(name.clone(), package.clone());
        }
    }

    // Process only the files to generate
    for file_name in &request.file_to_generate {
        let file_descriptor = request
            .proto_file
            .iter()
            .find(|f| f.name.as_ref() == Some(file_name))
            .context("File descriptor not found")?;

        let package_name = file_descriptor.package.clone().unwrap_or_default();

        // Extract Go package from options
        let go_package = if let Some(options) = &file_descriptor.options {
            if let Some(go_pkg) = &options.go_package {
                // go_package format: "github.com/org/repo/path"
                // Extract last component as package name
                go_pkg.split('/').next_back().unwrap_or("pb").to_string()
            } else {
                package_name
                    .split('.')
                    .next_back()
                    .unwrap_or("pb")
                    .to_string()
            }
        } else {
            package_name
                .split('.')
                .next_back()
                .unwrap_or("pb")
                .to_string()
        };

        let mut services_found = false;
        let mut service_definitions = String::new();

        // Generate code for each service
        for service in &file_descriptor.service {
            services_found = true;
            let service_name = service.name.clone().context("Service name missing")?;

            let mut client_methods = String::new();
            let mut client_method_impls = String::new();
            let mut server_methods = String::new();
            let mut unimplemented_methods = String::new();
            let mut handler_impls = String::new();
            let mut register_methods = String::new();

            // Generate methods
            for method in &service.method {
                let method_name = method.name.clone().context("Method name missing")?;
                let input_type = method
                    .input_type
                    .clone()
                    .context("Input type missing")?
                    .trim_start_matches('.')
                    .split('.')
                    .next_back()
                    .unwrap_or("")
                    .to_string();
                let output_type = method
                    .output_type
                    .clone()
                    .context("Output type missing")?
                    .trim_start_matches('.')
                    .split('.')
                    .next_back()
                    .unwrap_or("")
                    .to_string();

                let is_client_streaming = method.client_streaming.unwrap_or(false);
                let is_server_streaming = method.server_streaming.unwrap_or(false);

                // Determine RPC type and generate appropriate code
                match (is_client_streaming, is_server_streaming) {
                    (false, false) => {
                        // Unary-Unary
                        client_methods.push_str(
                            &UNARY_UNARY_CLIENT_METHOD
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type),
                        );

                        client_method_impls.push_str(
                            &UNARY_UNARY_CLIENT_IMPL
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type)
                                .replace("{{PACKAGE_NAME}}", &package_name),
                        );

                        server_methods.push_str(
                            &UNARY_UNARY_SERVER_METHOD
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type),
                        );

                        unimplemented_methods.push_str(
                            &UNIMPLEMENTED_METHOD
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type),
                        );

                        handler_impls.push_str(
                            &UNARY_UNARY_HANDLER_IMPL
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type),
                        );

                        register_methods.push_str(
                            &REGISTER_METHOD
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{PACKAGE_NAME}}", &package_name),
                        );
                    }
                    (false, true) => {
                        // Unary-Stream
                        client_methods.push_str(
                            &UNARY_STREAM_CLIENT_METHOD
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type)
                                .replace("{{SERVICE_NAME}}", &service_name),
                        );

                        client_method_impls.push_str(
                            &UNARY_STREAM_CLIENT_IMPL
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type)
                                .replace("{{PACKAGE_NAME}}", &package_name),
                        );

                        server_methods.push_str(
                            &UNARY_STREAM_SERVER_METHOD
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type)
                                .replace("{{SERVICE_NAME}}", &service_name),
                        );

                        unimplemented_methods.push_str(
                            &UNIMPLEMENTED_UNARY_STREAM_METHOD
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type),
                        );

                        handler_impls.push_str(
                            &UNARY_STREAM_HANDLER_IMPL
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type),
                        );

                        register_methods.push_str(
                            &REGISTER_UNARY_STREAM_METHOD
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{PACKAGE_NAME}}", &package_name),
                        );
                    }
                    (true, false) => {
                        // Stream-Unary
                        client_methods.push_str(
                            &STREAM_UNARY_CLIENT_METHOD
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type)
                                .replace("{{SERVICE_NAME}}", &service_name),
                        );

                        client_method_impls.push_str(
                            &STREAM_UNARY_CLIENT_IMPL
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type)
                                .replace("{{PACKAGE_NAME}}", &package_name),
                        );

                        server_methods.push_str(
                            &STREAM_UNARY_SERVER_METHOD
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type)
                                .replace("{{SERVICE_NAME}}", &service_name),
                        );

                        unimplemented_methods.push_str(
                            &UNIMPLEMENTED_STREAM_UNARY_METHOD
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type),
                        );

                        handler_impls.push_str(
                            &STREAM_UNARY_HANDLER_IMPL
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type),
                        );

                        register_methods.push_str(
                            &REGISTER_STREAM_UNARY_METHOD
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{PACKAGE_NAME}}", &package_name),
                        );
                    }
                    (true, true) => {
                        // Stream-Stream
                        client_methods.push_str(
                            &STREAM_STREAM_CLIENT_METHOD
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type)
                                .replace("{{SERVICE_NAME}}", &service_name),
                        );

                        client_method_impls.push_str(
                            &STREAM_STREAM_CLIENT_IMPL
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type)
                                .replace("{{PACKAGE_NAME}}", &package_name),
                        );

                        server_methods.push_str(
                            &STREAM_STREAM_SERVER_METHOD
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type)
                                .replace("{{SERVICE_NAME}}", &service_name),
                        );

                        unimplemented_methods.push_str(
                            &UNIMPLEMENTED_STREAM_STREAM_METHOD
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type),
                        );

                        handler_impls.push_str(
                            &STREAM_STREAM_HANDLER_IMPL
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{INPUT_TYPE}}", &input_type)
                                .replace("{{OUTPUT_TYPE}}", &output_type),
                        );

                        register_methods.push_str(
                            &REGISTER_STREAM_STREAM_METHOD
                                .replace("{{SERVICE_NAME}}", &service_name)
                                .replace("{{METHOD_NAME}}", &method_name)
                                .replace("{{PACKAGE_NAME}}", &package_name),
                        );
                    }
                }
            }

            // Generate client interface
            let client_interface = CLIENT_INTERFACE_TEMPLATE
                .replace("{{SERVICE_NAME}}", &service_name)
                .replace("{{CLIENT_METHODS}}", client_methods.trim_end())
                .replace("{{CLIENT_METHOD_IMPLS}}", &client_method_impls);
            service_definitions.push_str(&client_interface);

            // Generate server interface
            let server_interface = SERVER_INTERFACE_TEMPLATE
                .replace("{{SERVICE_NAME}}", &service_name)
                .replace("{{SERVER_METHODS}}", server_methods.trim_end())
                .replace("{{UNIMPLEMENTED_METHODS}}", &unimplemented_methods)
                .replace("{{HANDLER_IMPLS}}", &handler_impls)
                .replace("{{REGISTER_METHODS}}", register_methods.trim_end());
            service_definitions.push_str(&server_interface);
        }

        if services_found {
            // Generate the file content
            let proto_file_name = file_name.clone();
            let output_file_name = file_name
                .strip_suffix(".proto")
                .unwrap_or(file_name)
                .to_string()
                + "_slimrpc.pb.go";

            let content = FILE_TEMPLATE
                .replace("{{PROTO_FILE}}", &proto_file_name)
                .replace("{{PACKAGE}}", &go_package)
                .replace("{{PROTO_IMPORTS}}", "")
                .replace("{{SERVICE_DEFINITIONS}}", &service_definitions);

            response.file.push(File {
                name: Some(output_file_name),
                insertion_point: None,
                content: Some(content),
                generated_code_info: None,
            });
        }
    }

    Ok(response)
}

#[cfg(test)]
mod tests {
    use super::*;
    use prost_types::{
        DescriptorProto, FileDescriptorProto, MethodDescriptorProto, ServiceDescriptorProto,
    };

    fn create_test_file_descriptor(
        file_name: &str,
        package: &str,
        services: Vec<ServiceDescriptorProto>,
    ) -> FileDescriptorProto {
        FileDescriptorProto {
            name: Some(file_name.to_string()),
            package: Some(package.to_string()),
            service: services,
            message_type: vec![
                DescriptorProto {
                    name: Some("Request".to_string()),
                    ..Default::default()
                },
                DescriptorProto {
                    name: Some("Response".to_string()),
                    ..Default::default()
                },
            ],
            ..Default::default()
        }
    }

    #[test]
    fn test_generate_unary_unary() {
        let service = ServiceDescriptorProto {
            name: Some("TestService".to_string()),
            method: vec![MethodDescriptorProto {
                name: Some("UnaryCall".to_string()),
                input_type: Some(".test.Request".to_string()),
                output_type: Some(".test.Response".to_string()),
                client_streaming: Some(false),
                server_streaming: Some(false),
                ..Default::default()
            }],
            ..Default::default()
        };

        let file = create_test_file_descriptor("test.proto", "test", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["test.proto".to_string()],
            proto_file: vec![file],
            ..Default::default()
        };

        let response = generate(request).unwrap();
        assert_eq!(response.file.len(), 1);

        let content = response.file[0].content.as_ref().unwrap();
        assert!(content.contains("type TestServiceClient interface"));
        assert!(content.contains("type TestServiceServer interface"));
        assert!(
            content.contains("UnaryCall(ctx context.Context, req *Request) (*Response, error)")
        );
    }
}
