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

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

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

{{PACKAGE_DECL}}

import io.agntcy.slim.bindings.BidiStreamHandler;
import io.agntcy.slim.bindings.Channel;
import io.agntcy.slim.bindings.Context;
import io.agntcy.slim.bindings.RequestStream;
import io.agntcy.slim.bindings.RequestStreamWriter;
import io.agntcy.slim.bindings.ResponseSink;
import io.agntcy.slim.bindings.ResponseStreamReader;
import io.agntcy.slim.bindings.RpcCode;
import io.agntcy.slim.bindings.RpcException;
import io.agntcy.slim.bindings.Server;
import io.agntcy.slim.bindings.UnaryStreamHandler;
import io.agntcy.slim.bindings.UnaryUnaryHandler;
import io.agntcy.slim.bindings.StreamUnaryHandler;
import io.agntcy.slim.bindings.StreamStreamHandler;
import io.agntcy.slim.bindings.slimrpc.ClientRequestStream;
import io.agntcy.slim.bindings.slimrpc.ClientBidiStream;

import static io.agntcy.slim.bindings.slimrpc.SlimrpcUtil.parse;
import static io.agntcy.slim.bindings.slimrpc.SlimrpcUtil.toRpcException;

import java.time.Duration;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

{{SERVICE_DEFINITION}}
"#;

const SERVICE_TEMPLATE: &str = r#"public final class {{SERVICE_NAME}}Slimrpc {
    private {{SERVICE_NAME}}Slimrpc() {}

    public static final String SERVICE_NAME = "{{FULL_SERVICE_NAME}}";

    public interface {{SERVICE_NAME}}Client {
{{CLIENT_METHODS}}
    }

    public static final class {{SERVICE_NAME}}ClientImpl implements {{SERVICE_NAME}}Client {
        private final Channel channel;

        public {{SERVICE_NAME}}ClientImpl(Channel channel) {
            this.channel = channel;
        }

{{CLIENT_METHOD_IMPLS}}
    }

    public interface {{SERVICE_NAME}}Server {
{{SERVER_METHODS}}
    }

    public abstract static class Unimplemented{{SERVICE_NAME}}Server implements {{SERVICE_NAME}}Server {
{{UNIMPLEMENTED_METHODS}}
    }

    public static void register{{SERVICE_NAME}}Server(Server server, {{SERVICE_NAME}}Server impl) {
{{REGISTER_METHODS}}
    }

{{HANDLER_IMPLS}}
}
"#;

const CLIENT_UNARY_UNARY_METHOD: &str = r#"        {{OUTPUT_TYPE}} {{METHOD_NAME}}({{INPUT_TYPE}} request, Duration timeout, Map<String, String> metadata) throws RpcException;
"#;

const CLIENT_UNARY_STREAM_METHOD: &str = r#"        ResponseStreamReader {{METHOD_NAME}}({{INPUT_TYPE}} request, Duration timeout, Map<String, String> metadata) throws RpcException;
"#;

const CLIENT_STREAM_UNARY_METHOD: &str = r#"        ClientRequestStream<{{INPUT_TYPE}}, {{OUTPUT_TYPE}}> {{METHOD_NAME}}(Duration timeout, Map<String, String> metadata) throws RpcException;
"#;

const CLIENT_STREAM_STREAM_METHOD: &str = r#"        ClientBidiStream<{{INPUT_TYPE}}> {{METHOD_NAME}}(Duration timeout, Map<String, String> metadata) throws RpcException;
"#;

const CLIENT_UNARY_UNARY_IMPL: &str = r#"        @Override
        public {{OUTPUT_TYPE}} {{METHOD_NAME}}({{INPUT_TYPE}} request, Duration timeout, Map<String, String> metadata) throws RpcException {
            byte[] bytes = channel.callUnaryAsync(
                SERVICE_NAME,
                "{{METHOD_NAME}}",
                request.toByteArray(),
                timeout,
                metadata
            ).join();
            return parse(bytes, {{OUTPUT_TYPE}}.parser());
        }

"#;

const CLIENT_UNARY_STREAM_IMPL: &str = r#"        @Override
        public ResponseStreamReader {{METHOD_NAME}}({{INPUT_TYPE}} request, Duration timeout, Map<String, String> metadata) throws RpcException {
            return channel.callUnaryStreamAsync(
                SERVICE_NAME,
                "{{METHOD_NAME}}",
                request.toByteArray(),
                timeout,
                metadata
            ).join();
        }

"#;

const CLIENT_STREAM_UNARY_IMPL: &str = r#"        @Override
        public ClientRequestStream<{{INPUT_TYPE}}, {{OUTPUT_TYPE}}> {{METHOD_NAME}}(Duration timeout, Map<String, String> metadata) throws RpcException {
            RequestStreamWriter writer = channel.callStreamUnary(
                    SERVICE_NAME,
                    "{{METHOD_NAME}}",
                    timeout,
                    metadata
            );
            return new ClientRequestStream<>(writer, {{INPUT_TYPE}}::toByteArray, bytes -> parse(bytes, {{OUTPUT_TYPE}}.parser()));
        }

"#;

const CLIENT_STREAM_STREAM_IMPL: &str = r#"        @Override
        public ClientBidiStream<{{INPUT_TYPE}}> {{METHOD_NAME}}(Duration timeout, Map<String, String> metadata) throws RpcException {
            BidiStreamHandler handler = channel.callStreamStream(
                    SERVICE_NAME,
                    "{{METHOD_NAME}}",
                    timeout,
                    metadata
            );
            return new ClientBidiStream<>(handler, {{INPUT_TYPE}}::toByteArray);
        }

"#;

const SERVER_UNARY_UNARY_METHOD: &str = r#"        CompletableFuture<{{OUTPUT_TYPE}}> {{METHOD_NAME}}({{INPUT_TYPE}} request, Context context);
"#;

const SERVER_UNARY_STREAM_METHOD: &str = r#"        CompletableFuture<Void> {{METHOD_NAME}}({{INPUT_TYPE}} request, Context context, ResponseSink sink);
"#;

const SERVER_STREAM_UNARY_METHOD: &str = r#"        CompletableFuture<{{OUTPUT_TYPE}}> {{METHOD_NAME}}(RequestStream stream, Context context);
"#;

const SERVER_STREAM_STREAM_METHOD: &str = r#"        CompletableFuture<Void> {{METHOD_NAME}}(RequestStream stream, Context context, ResponseSink sink);
"#;

const UNIMPLEMENTED_UNARY_UNARY: &str = r#"        @Override
        public CompletableFuture<{{OUTPUT_TYPE}}> {{METHOD_NAME}}({{INPUT_TYPE}} request, Context context) {
            return CompletableFuture.failedFuture(new UnsupportedOperationException("Method {{METHOD_NAME}} not implemented"));
        }

"#;

const UNIMPLEMENTED_UNARY_STREAM: &str = r#"        @Override
        public CompletableFuture<Void> {{METHOD_NAME}}({{INPUT_TYPE}} request, Context context, ResponseSink sink) {
            return CompletableFuture.failedFuture(new UnsupportedOperationException("Method {{METHOD_NAME}} not implemented"));
        }

"#;

const UNIMPLEMENTED_STREAM_UNARY: &str = r#"        @Override
        public CompletableFuture<{{OUTPUT_TYPE}}> {{METHOD_NAME}}(RequestStream stream, Context context) {
            return CompletableFuture.failedFuture(new UnsupportedOperationException("Method {{METHOD_NAME}} not implemented"));
        }

"#;

const UNIMPLEMENTED_STREAM_STREAM: &str = r#"        @Override
        public CompletableFuture<Void> {{METHOD_NAME}}(RequestStream stream, Context context, ResponseSink sink) {
            return CompletableFuture.failedFuture(new UnsupportedOperationException("Method {{METHOD_NAME}} not implemented"));
        }

"#;

const REGISTER_UNARY_UNARY: &str = r#"        server.registerUnaryUnary(SERVICE_NAME, "{{METHOD_NAME}}", new {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler(impl));
"#;

const REGISTER_UNARY_STREAM: &str = r#"        server.registerUnaryStream(SERVICE_NAME, "{{METHOD_NAME}}", new {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler(impl));
"#;

const REGISTER_STREAM_UNARY: &str = r#"        server.registerStreamUnary(SERVICE_NAME, "{{METHOD_NAME}}", new {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler(impl));
"#;

const REGISTER_STREAM_STREAM: &str = r#"        server.registerStreamStream(SERVICE_NAME, "{{METHOD_NAME}}", new {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler(impl));
"#;

const HANDLER_UNARY_UNARY: &str = r#"    private static final class {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler implements UnaryUnaryHandler {
        private final {{SERVICE_NAME}}Server impl;

        private {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler({{SERVICE_NAME}}Server impl) {
            this.impl = impl;
        }

        @Override
        public CompletableFuture<byte[]> handle(byte[] request, Context context) {
            try {
                {{INPUT_TYPE}} parsed = parse(request, {{INPUT_TYPE}}.parser());
                return impl.{{METHOD_NAME}}(parsed, context)
                        .thenApply(resp -> resp.toByteArray());
            } catch (Exception e) {
                return CompletableFuture.failedFuture(e);
            }
        }
    }

"#;

const HANDLER_UNARY_STREAM: &str = r#"    private static final class {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler implements UnaryStreamHandler {
        private final {{SERVICE_NAME}}Server impl;

        private {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler({{SERVICE_NAME}}Server impl) {
            this.impl = impl;
        }

        @Override
        public CompletableFuture<Void> handle(byte[] request, Context context, ResponseSink sink) {
            try {
                {{INPUT_TYPE}} parsed = parse(request, {{INPUT_TYPE}}.parser());
                return impl.{{METHOD_NAME}}(parsed, context, sink)
                        .handle((ignored, error) -> {
                            if (error != null) {
                                RpcException rpcErr = toRpcException(error, RpcCode.INTERNAL);
                                return (java.util.concurrent.CompletionStage<Void>) sink.sendErrorAsync(rpcErr)
                                    .thenCompose(v -> CompletableFuture.<Void>failedFuture(rpcErr));
                            }
                            return (java.util.concurrent.CompletionStage<Void>) sink.closeAsync();
                        })
                        .thenCompose((java.util.concurrent.CompletionStage<Void> stage) -> stage);
            } catch (Exception e) {
                RpcException rpcErr = new RpcException.Rpc(
                        RpcCode.INVALID_ARGUMENT,
                        e.getMessage() == null ? "invalid request" : e.getMessage(),
                        null
                );
                return sink.sendErrorAsync(rpcErr)
                    .thenCompose(v -> CompletableFuture.<Void>failedFuture(rpcErr));
            }
        }
    }

"#;

const HANDLER_STREAM_UNARY: &str = r#"    private static final class {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler implements StreamUnaryHandler {
        private final {{SERVICE_NAME}}Server impl;

        private {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler({{SERVICE_NAME}}Server impl) {
            this.impl = impl;
        }

        @Override
        public CompletableFuture<byte[]> handle(RequestStream stream, Context context) {
            return impl.{{METHOD_NAME}}(stream, context)
                    .handle((resp, error) -> {
                        if (error != null) {
                            return (java.util.concurrent.CompletionStage<byte[]>) CompletableFuture.<byte[]>failedFuture(
                                toRpcException(error, RpcCode.INTERNAL)
                            );
                        }
                        return (java.util.concurrent.CompletionStage<byte[]>) CompletableFuture.completedFuture(
                                resp.toByteArray()
                        );
                    })
                    .thenCompose((java.util.concurrent.CompletionStage<byte[]> stage) -> stage);
        }
    }

"#;

const HANDLER_STREAM_STREAM: &str = r#"    private static final class {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler implements StreamStreamHandler {
        private final {{SERVICE_NAME}}Server impl;

        private {{SERVICE_NAME}}_{{METHOD_NAME}}_Handler({{SERVICE_NAME}}Server impl) {
            this.impl = impl;
        }

        @Override
        public CompletableFuture<Void> handle(RequestStream stream, Context context, ResponseSink sink) {
            return impl.{{METHOD_NAME}}(stream, context, sink)
                    .handle((ignored, error) -> {
                        if (error != null) {
                            RpcException rpcErr = toRpcException(error, RpcCode.INTERNAL);
                            return (java.util.concurrent.CompletionStage<Void>) sink.sendErrorAsync(rpcErr)
                                    .thenCompose(v -> CompletableFuture.<Void>failedFuture(rpcErr));
                        }
                        return (java.util.concurrent.CompletionStage<Void>) sink.closeAsync();
                    })
                    .thenCompose((java.util.concurrent.CompletionStage<Void> stage) -> stage);
        }
    }

"#;

struct RpcTemplates {
    client_method: &'static str,
    client_impl: &'static str,
    server_method: &'static str,
    unimplemented: &'static str,
    register: &'static str,
    handler: &'static str,
}

const UNARY_UNARY_TEMPLATES: RpcTemplates = RpcTemplates {
    client_method: CLIENT_UNARY_UNARY_METHOD,
    client_impl: CLIENT_UNARY_UNARY_IMPL,
    server_method: SERVER_UNARY_UNARY_METHOD,
    unimplemented: UNIMPLEMENTED_UNARY_UNARY,
    register: REGISTER_UNARY_UNARY,
    handler: HANDLER_UNARY_UNARY,
};

const UNARY_STREAM_TEMPLATES: RpcTemplates = RpcTemplates {
    client_method: CLIENT_UNARY_STREAM_METHOD,
    client_impl: CLIENT_UNARY_STREAM_IMPL,
    server_method: SERVER_UNARY_STREAM_METHOD,
    unimplemented: UNIMPLEMENTED_UNARY_STREAM,
    register: REGISTER_UNARY_STREAM,
    handler: HANDLER_UNARY_STREAM,
};

const STREAM_UNARY_TEMPLATES: RpcTemplates = RpcTemplates {
    client_method: CLIENT_STREAM_UNARY_METHOD,
    client_impl: CLIENT_STREAM_UNARY_IMPL,
    server_method: SERVER_STREAM_UNARY_METHOD,
    unimplemented: UNIMPLEMENTED_STREAM_UNARY,
    register: REGISTER_STREAM_UNARY,
    handler: HANDLER_STREAM_UNARY,
};

const STREAM_STREAM_TEMPLATES: RpcTemplates = RpcTemplates {
    client_method: CLIENT_STREAM_STREAM_METHOD,
    client_impl: CLIENT_STREAM_STREAM_IMPL,
    server_method: SERVER_STREAM_STREAM_METHOD,
    unimplemented: UNIMPLEMENTED_STREAM_STREAM,
    register: REGISTER_STREAM_STREAM,
    handler: HANDLER_STREAM_STREAM,
};

fn rpc_templates(client_streaming: bool, server_streaming: bool) -> &'static RpcTemplates {
    match (client_streaming, server_streaming) {
        (false, false) => &UNARY_UNARY_TEMPLATES,
        (false, true) => &UNARY_STREAM_TEMPLATES,
        (true, false) => &STREAM_UNARY_TEMPLATES,
        (true, true) => &STREAM_STREAM_TEMPLATES,
    }
}

fn apply(template: &str, vars: &[(&str, &str)]) -> String {
    let mut result = template.to_string();
    for (key, value) in vars {
        result = result.replace(key, value);
    }
    result
}

#[derive(Clone)]
struct JavaFileInfo {
    java_package: String,
    outer_class: String,
    multiple_files: bool,
}

fn file_name_base(file_name: &str) -> String {
    file_name
        .split('/')
        .next_back()
        .unwrap_or(file_name)
        .strip_suffix(".proto")
        .unwrap_or(file_name)
        .to_string()
}

fn java_file_info(file: &FileDescriptorProto) -> JavaFileInfo {
    let java_package = file
        .options
        .as_ref()
        .and_then(|opts| opts.java_package.clone())
        .filter(|s: &String| !s.is_empty())
        .or_else(|| file.package.clone())
        .unwrap_or_default();

    let base_name = file
        .name
        .as_deref()
        .map(file_name_base)
        .unwrap_or_else(|| "unknown_proto".to_string());

    let outer_class = file
        .options
        .as_ref()
        .and_then(|opts| opts.java_outer_classname.clone())
        .filter(|s: &String| !s.is_empty())
        .unwrap_or_else(|| format!("{}Proto", base_name.to_upper_camel_case()));

    let multiple_files = file
        .options
        .as_ref()
        .and_then(|opts| opts.java_multiple_files)
        .unwrap_or(false);

    JavaFileInfo {
        java_package,
        outer_class,
        multiple_files,
    }
}

fn collect_message_types(
    proto_prefix: String,
    java_prefix: String,
    messages: &[DescriptorProto],
    type_map: &mut HashMap<String, String>,
) {
    for msg in messages {
        let Some(name) = msg.name.as_ref() else {
            continue;
        };

        let proto_full = if proto_prefix.is_empty() {
            format!(".{}", name)
        } else {
            format!(".{}.{}", proto_prefix, name)
        };

        let java_full = if java_prefix.is_empty() {
            name.to_string()
        } else {
            format!("{}.{}", java_prefix, name)
        };

        type_map.insert(proto_full.clone(), java_full.clone());

        let nested_proto_prefix = if proto_prefix.is_empty() {
            name.to_string()
        } else {
            format!("{}.{}", proto_prefix, name)
        };

        collect_message_types(nested_proto_prefix, java_full, &msg.nested_type, type_map);
    }
}

fn resolve_java_type(
    full_type_name: &str,
    current_java_package: &str,
    type_map: &HashMap<String, String>,
) -> String {
    let trimmed = full_type_name.trim_start_matches('.');

    if trimmed.starts_with("google.protobuf.") {
        let type_name = trimmed.rsplit('.').next().unwrap_or(trimmed);
        return format!("com.google.protobuf.{}", type_name);
    }

    let lookup_key = format!(".{}", trimmed);
    if let Some(java_full) = type_map.get(&lookup_key) {
        if !current_java_package.is_empty() {
            let prefix = format!("{}.", current_java_package);
            if java_full.starts_with(&prefix) {
                return java_full.trim_start_matches(&prefix).to_string();
            }
        }
        return java_full.clone();
    }

    trimmed.rsplit('.').next().unwrap_or(trimmed).to_string()
}

/// Generate Java slimrpc code from a CodeGeneratorRequest
pub fn generate(request: CodeGeneratorRequest) -> Result<CodeGeneratorResponse> {
    let mut response = CodeGeneratorResponse::default();

    let mut file_map: HashMap<String, FileDescriptorProto> = HashMap::new();
    let mut file_info_map: HashMap<String, JavaFileInfo> = HashMap::new();

    for file in &request.proto_file {
        if let Some(name) = &file.name {
            file_map.insert(name.clone(), file.clone());
            file_info_map.insert(name.clone(), java_file_info(file));
        }
    }

    let mut type_map: HashMap<String, String> = HashMap::new();

    for file in &request.proto_file {
        let Some(file_name) = file.name.as_ref() else {
            continue;
        };
        let info = file_info_map
            .get(file_name.as_str())
            .cloned()
            .unwrap_or_else(|| java_file_info(file));

        let java_prefix = if info.multiple_files {
            info.java_package.clone()
        } else if info.java_package.is_empty() {
            info.outer_class.clone()
        } else {
            format!("{}.{}", info.java_package, info.outer_class)
        };

        let proto_prefix = file.package.clone().unwrap_or_default();
        collect_message_types(proto_prefix, java_prefix, &file.message_type, &mut type_map);
    }

    for file_name in &request.file_to_generate {
        let file_descriptor = file_map
            .remove(file_name)
            .context("File descriptor not found")?;

        let info = file_info_map
            .remove(file_name)
            .unwrap_or_else(|| java_file_info(&file_descriptor));

        let java_package = info.java_package;
        let package_decl = if java_package.is_empty() {
            String::new()
        } else {
            format!("package {};", java_package)
        };

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

        for service in file_descriptor.service {
            let service_name = service.name.context("Service name missing")?;
            let full_service_name = if package_name.is_empty() {
                service_name.clone()
            } else {
                format!("{}.{}", package_name, service_name)
            };

            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 register_methods = String::new();
            let mut handler_impls = String::new();

            for method in service.method {
                let method_name = method.name.as_deref().context("Method name missing")?;
                let input_type = method.input_type.as_deref().context("Input type missing")?;
                let output_type = method
                    .output_type
                    .as_deref()
                    .context("Output type missing")?;

                let input_java = resolve_java_type(input_type, &java_package, &type_map);
                let output_java = resolve_java_type(output_type, &java_package, &type_map);

                let t = rpc_templates(
                    method.client_streaming.unwrap_or(false),
                    method.server_streaming.unwrap_or(false),
                );
                let vars = &[
                    ("{{METHOD_NAME}}", method_name),
                    ("{{SERVICE_NAME}}", service_name.as_str()),
                    ("{{INPUT_TYPE}}", input_java.as_str()),
                    ("{{OUTPUT_TYPE}}", output_java.as_str()),
                ];

                client_methods.push_str(&apply(t.client_method, vars));
                client_method_impls.push_str(&apply(t.client_impl, vars));
                server_methods.push_str(&apply(t.server_method, vars));
                unimplemented_methods.push_str(&apply(t.unimplemented, vars));
                register_methods.push_str(&apply(t.register, vars));
                handler_impls.push_str(&apply(t.handler, vars));
            }

            let service_content = SERVICE_TEMPLATE
                .replace("{{SERVICE_NAME}}", &service_name)
                .replace("{{FULL_SERVICE_NAME}}", &full_service_name)
                .replace("{{CLIENT_METHODS}}", &client_methods)
                .replace("{{CLIENT_METHOD_IMPLS}}", &client_method_impls)
                .replace("{{SERVER_METHODS}}", &server_methods)
                .replace("{{UNIMPLEMENTED_METHODS}}", &unimplemented_methods)
                .replace("{{REGISTER_METHODS}}", register_methods.trim_end())
                .replace("{{HANDLER_IMPLS}}", &handler_impls);

            let output_file_name = format!("{}Slimrpc.java", service_name);
            let output_path = if java_package.is_empty() {
                output_file_name
            } else {
                format!("{}/{}", java_package.replace('.', "/"), output_file_name)
            };

            let final_content = FILE_TEMPLATE
                .replace("{{PROTO_FILE}}", file_name)
                .replace("{{PACKAGE_DECL}}", &package_decl)
                .replace("{{SERVICE_DEFINITION}}", &service_content);

            response.file.push(File {
                name: Some(output_path),
                content: Some(final_content),
                ..Default::default()
            });
        }
    }

    Ok(response)
}

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

    fn create_test_file(
        name: &str,
        package: &str,
        services: Vec<ServiceDescriptorProto>,
        messages: Vec<DescriptorProto>,
    ) -> FileDescriptorProto {
        FileDescriptorProto {
            name: Some(name.to_string()),
            package: Some(package.to_string()),
            service: services,
            message_type: messages,
            ..Default::default()
        }
    }

    fn create_method(
        name: &str,
        input: &str,
        output: &str,
        client_streaming: bool,
        server_streaming: bool,
    ) -> MethodDescriptorProto {
        MethodDescriptorProto {
            name: Some(name.to_string()),
            input_type: Some(input.to_string()),
            output_type: Some(output.to_string()),
            client_streaming: Some(client_streaming),
            server_streaming: Some(server_streaming),
            ..Default::default()
        }
    }

    fn create_service(name: &str, methods: Vec<MethodDescriptorProto>) -> ServiceDescriptorProto {
        ServiceDescriptorProto {
            name: Some(name.to_string()),
            method: methods,
            ..Default::default()
        }
    }

    fn create_message(name: &str) -> DescriptorProto {
        DescriptorProto {
            name: Some(name.to_string()),
            ..Default::default()
        }
    }

    #[test]
    fn test_generate_no_services() {
        let file = create_test_file("test.proto", "test", vec![], vec![]);
        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(), 0);
    }

    #[test]
    fn test_generate_unary_unary() {
        let method = create_method(
            "GetItem",
            ".test.GetItemRequest",
            ".test.GetItemResponse",
            false,
            false,
        );
        let service = create_service("ItemService", vec![method]);
        let messages = vec![
            create_message("GetItemRequest"),
            create_message("GetItemResponse"),
        ];
        let file = create_test_file("test.proto", "test", vec![service], messages);

        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("public interface ItemServiceClient"));
        assert!(content.contains("public interface ItemServiceServer"));
        assert!(
            content.contains("TestProto.GetItemResponse GetItem(TestProto.GetItemRequest request,")
        );
        assert!(content.contains(
            "CompletableFuture<TestProto.GetItemResponse> GetItem(TestProto.GetItemRequest request, Context context)"
        ));
        assert!(content.contains("registerUnaryUnary"));
        assert!(content.contains("implements UnaryUnaryHandler"));
    }

    #[test]
    fn test_generate_all_streaming_types() {
        let methods = vec![
            create_method("Unary", ".pkg.Req", ".pkg.Res", false, false),
            create_method("ServerStream", ".pkg.Req", ".pkg.Res", false, true),
            create_method("ClientStream", ".pkg.Req", ".pkg.Res", true, false),
            create_method("BidiStream", ".pkg.Req", ".pkg.Res", true, true),
        ];
        let service = create_service("TestService", methods);
        let messages = vec![create_message("Req"), create_message("Res")];
        let file = create_test_file("test.proto", "pkg", vec![service], messages);

        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("registerUnaryUnary"));
        assert!(content.contains("implements UnaryUnaryHandler"));
        assert!(content.contains("registerUnaryStream"));
        assert!(content.contains("implements UnaryStreamHandler"));
        assert!(content.contains("ResponseStreamReader ServerStream(TestProto.Req request,"));
        assert!(content.contains("registerStreamUnary"));
        assert!(content.contains("implements StreamUnaryHandler"));
        assert!(
            content.contains("ClientRequestStream<TestProto.Req, TestProto.Res> ClientStream(")
        );
        assert!(content.contains("registerStreamStream"));
        assert!(content.contains("implements StreamStreamHandler"));
        assert!(content.contains("ClientBidiStream<TestProto.Req> BidiStream("));
    }

    #[test]
    fn test_generate_multiple_services() {
        let service1 = create_service(
            "ServiceA",
            vec![create_method(
                "MethodA", ".pkg.Req", ".pkg.Res", false, false,
            )],
        );
        let service2 = create_service(
            "ServiceB",
            vec![create_method(
                "MethodB", ".pkg.Req", ".pkg.Res", false, false,
            )],
        );
        let messages = vec![create_message("Req"), create_message("Res")];
        let file = create_test_file("test.proto", "pkg", vec![service1, service2], messages);

        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(), 2);

        let content_a = response.file[0].content.as_ref().unwrap();
        assert!(content_a.contains("class ServiceASlimrpc"));
        assert!(content_a.contains("ServiceAClient"));

        let content_b = response.file[1].content.as_ref().unwrap();
        assert!(content_b.contains("class ServiceBSlimrpc"));
        assert!(content_b.contains("ServiceBClient"));
    }

    #[test]
    fn test_generate_empty_package() {
        let method = create_method("Call", ".Request", ".Response", false, false);
        let service = create_service("MyService", vec![method]);
        let messages = vec![create_message("Request"), create_message("Response")];
        let file = create_test_file("test.proto", "", vec![service], messages);

        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("package "));
        assert!(content.contains("SERVICE_NAME = \"MyService\""));
        assert_eq!(
            response.file[0].name.as_ref().unwrap(),
            "MyServiceSlimrpc.java"
        );
    }

    #[test]
    fn test_generate_google_protobuf_type() {
        let method = create_method(
            "Delete",
            ".test.DeleteRequest",
            ".google.protobuf.Empty",
            false,
            false,
        );
        let service = create_service("TestService", vec![method]);
        let messages = vec![create_message("DeleteRequest")];
        let file = create_test_file("test.proto", "test", vec![service], messages);

        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("com.google.protobuf.Empty"));
    }

    #[test]
    fn test_generate_nested_message_types() {
        let inner = DescriptorProto {
            name: Some("Inner".to_string()),
            ..Default::default()
        };
        let outer = DescriptorProto {
            name: Some("Outer".to_string()),
            nested_type: vec![inner],
            ..Default::default()
        };
        let method = create_method(
            "DoSomething",
            ".test.Outer.Inner",
            ".test.Outer",
            false,
            false,
        );
        let service = create_service("TestService", vec![method]);
        let file = FileDescriptorProto {
            name: Some("test.proto".to_string()),
            package: Some("test".to_string()),
            service: vec![service],
            message_type: vec![outer],
            ..Default::default()
        };

        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("TestProto.Outer.Inner request"),
            "content was:\n{}",
            content
        );
        assert!(
            content.contains("CompletableFuture<TestProto.Outer>"),
            "content was:\n{}",
            content
        );
    }

    #[test]
    fn test_generate_java_multiple_files() {
        let method = create_method("GetItem", ".test.Request", ".test.Response", false, false);
        let service = create_service("ItemService", vec![method]);
        let messages = vec![create_message("Request"), create_message("Response")];
        let file = FileDescriptorProto {
            name: Some("test.proto".to_string()),
            package: Some("test".to_string()),
            service: vec![service],
            message_type: messages,
            options: Some(FileOptions {
                java_multiple_files: Some(true),
                java_package: Some("com.example.test".to_string()),
                ..Default::default()
            }),
            ..Default::default()
        };

        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("Request request"),
            "content was:\n{}",
            content
        );
        assert!(
            content.contains("CompletableFuture<Response>"),
            "content was:\n{}",
            content
        );
        let path = response.file[0].name.as_ref().unwrap();
        assert!(path.starts_with("com/example/test/"), "path was: {}", path);
    }

    #[test]
    fn test_generate_java_package_option() {
        let method = create_method("Call", ".test.Request", ".test.Response", false, false);
        let service = create_service("TestService", vec![method]);
        let messages = vec![create_message("Request"), create_message("Response")];
        let file = FileDescriptorProto {
            name: Some("test.proto".to_string()),
            package: Some("test".to_string()),
            service: vec![service],
            message_type: messages,
            options: Some(FileOptions {
                java_package: Some("com.example.myapp".to_string()),
                ..Default::default()
            }),
            ..Default::default()
        };

        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("package com.example.myapp;"));

        let path = response.file[0].name.as_ref().unwrap();
        assert!(path.starts_with("com/example/myapp/"), "path was: {}", path);
    }
}
