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

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

/// Supported parameters for the Python code generator plugin.
const TYPES_IMPORT: &str = "types_import";

// --- TEMPLATE DEFINITIONS ---

const FILE_TEMPLATE: &str = r#"# Generated by protoc-slimrpc-plugin
# This file is generated from an embedded template. Do not edit directly.

from datetime import timedelta
from typing import Optional

from google.rpc import code_pb2 as code__pb2

import slim_bindings

{{PB2_IMPORTS}}

{{SERVICE_DEFINITIONS}}
{{HANDLER_CLASSES}}
{{ADD_SERVICER_FUNCTIONS}}"#;

const SERVICE_STUB_TEMPLATE: &str = r#"
class {{SERVICE_NAME}}Stub:
    """Client stub for {{SERVICE_NAME}}."""

    def __init__(self, channel):
        """Constructor.

        Args:
            channel: A slim_bindings.Channel.
        """
        self._channel = channel

{{METHOD_STUB_METHODS}}
"#;

const UNARY_UNARY_STUB_METHOD_TEMPLATE: &str = r#"    async def {{METHOD_NAME}}(self, request: {{INPUT_TYPE_FULL_PATH}}, timeout: Optional[timedelta] = None, metadata: Optional[dict[str, str]] = None) -> {{OUTPUT_TYPE_FULL_PATH}}:
        """Call {{METHOD_NAME}} method."""
        response_bytes = await self._channel.call_unary_async(
            "{{PACKAGE_NAME}}.{{SERVICE_NAME}}",
            "{{METHOD_NAME}}",
            {{INPUT_TYPE_FULL_PATH}}.SerializeToString(request),
            timeout,
            metadata,
        )
        return {{OUTPUT_TYPE_FULL_PATH}}.FromString(response_bytes)

"#;

const UNARY_STREAM_STUB_METHOD_TEMPLATE: &str = r#"    async def {{METHOD_NAME}}(self, request: {{INPUT_TYPE_FULL_PATH}}, timeout: Optional[timedelta] = None, metadata: Optional[dict[str, str]] = None):
        """Call {{METHOD_NAME}} method."""
        response_stream = await self._channel.call_unary_stream_async(
            "{{PACKAGE_NAME}}.{{SERVICE_NAME}}",
            "{{METHOD_NAME}}",
            {{INPUT_TYPE_FULL_PATH}}.SerializeToString(request),
            timeout,
            metadata,
        )
        while True:
            stream_msg = await response_stream.next_async()
            if stream_msg.is_end():
                break
            if stream_msg.is_error():
                raise stream_msg[0]
            if stream_msg.is_data():
                yield {{OUTPUT_TYPE_FULL_PATH}}.FromString(stream_msg[0])

"#;

const STREAM_UNARY_STUB_METHOD_TEMPLATE: &str = r#"    async def {{METHOD_NAME}}(self, request_iterator, timeout: Optional[timedelta] = None, metadata: Optional[dict[str, str]] = None) -> {{OUTPUT_TYPE_FULL_PATH}}:
        """Call {{METHOD_NAME}} method."""
        request_stream = self._channel.call_stream_unary(
            "{{PACKAGE_NAME}}.{{SERVICE_NAME}}",
            "{{METHOD_NAME}}",
            timeout,
            metadata,
        )
        async for request in request_iterator:
            await request_stream.send_async({{INPUT_TYPE_FULL_PATH}}.SerializeToString(request))
        response_bytes = await request_stream.finalize_async()
        return {{OUTPUT_TYPE_FULL_PATH}}.FromString(response_bytes)

"#;

const STREAM_STREAM_STUB_METHOD_TEMPLATE: &str = r#"    async def {{METHOD_NAME}}(self, request_iterator, timeout: Optional[timedelta] = None, metadata: Optional[dict[str, str]] = None):
        """Call {{METHOD_NAME}} method."""
        bidi_stream = self._channel.call_stream_stream(
            "{{PACKAGE_NAME}}.{{SERVICE_NAME}}",
            "{{METHOD_NAME}}",
            timeout,
            metadata,
        )

        async def send_requests():
            async for request in request_iterator:
                await bidi_stream.send_async({{INPUT_TYPE_FULL_PATH}}.SerializeToString(request))
            await bidi_stream.close_send_async()

        async def receive_responses():
            while True:
                stream_msg = await bidi_stream.recv_async()
                if stream_msg.is_end():
                    break
                if stream_msg.is_error():
                    raise stream_msg[0]
                if stream_msg.is_data():
                    yield {{OUTPUT_TYPE_FULL_PATH}}.FromString(stream_msg[0])

        # Start sending in background
        import asyncio
        send_task = asyncio.create_task(send_requests())

        try:
            async for response in receive_responses():
                yield response
        finally:
            await send_task

"#;

const SERVICE_SERVICER_TEMPLATE: &str = r#"
class {{SERVICE_NAME}}Servicer:
    """Server servicer for {{SERVICE_NAME}}. Implement this class to provide your service logic."""

{{METHOD_SERVICERS}}"#;

const STREAM_METHOD_SERVICER_TEMPLATE: &str = r#"    def {{METHOD_NAME}}(self, request_iterator, context):
        """Method for {{METHOD_NAME}}. Implement your service logic here."""
        raise slim_bindings.RpcError.Rpc(
            code=slim_bindings.RpcCode.UNIMPLEMENTED,
            message="Method not implemented!",
            details=None
        )

"#;

const UNARY_METHOD_SERVICER_TEMPLATE: &str = r#"    def {{METHOD_NAME}}(self, request, context):
        """Method for {{METHOD_NAME}}. Implement your service logic here."""
        raise slim_bindings.RpcError.Rpc(
            code=slim_bindings.RpcCode.UNIMPLEMENTED,
            message="Method not implemented!",
            details=None
        )

"#;

const ADD_SERVICER_FUNCTION_TEMPLATE: &str = r#"

def add_{{SERVICE_NAME}}Servicer_to_server(servicer, server: slim_bindings.Server):
{{REGISTER_METHODS}}
"#;

const UNARY_UNARY_HANDLER_CLASS_TEMPLATE: &str = r#"

class _{{SERVICE_NAME}}Servicer_{{METHOD_NAME}}_Handler:
    def __init__(self, servicer):
        self.servicer = servicer

    async def handle(self, request: bytes, context: slim_bindings.Context) -> bytes:
        try:
            request_msg = {{INPUT_TYPE_FULL_PATH}}.FromString(request)
            response = await self.servicer.{{METHOD_NAME}}(request_msg, context)
            return {{OUTPUT_TYPE_FULL_PATH}}.SerializeToString(response)
        except slim_bindings.RpcError:
            raise
        except Exception as e:
            raise slim_bindings.RpcError.Rpc(
                code=slim_bindings.RpcCode.INTERNAL,
                message=str(e),
                details=None
            )"#;

const UNARY_STREAM_HANDLER_CLASS_TEMPLATE: &str = r#"

class _{{SERVICE_NAME}}Servicer_{{METHOD_NAME}}_Handler:
    def __init__(self, servicer):
        self.servicer = servicer

    async def handle(self, request: bytes, context: slim_bindings.Context, sink: slim_bindings.ResponseSink):
        try:
            request_msg = {{INPUT_TYPE_FULL_PATH}}.FromString(request)
            response_iter = self.servicer.{{METHOD_NAME}}(request_msg, context)
            async for response in response_iter:
                await sink.send_async({{OUTPUT_TYPE_FULL_PATH}}.SerializeToString(response))
            await sink.close_async()
        except slim_bindings.RpcError as e:
            await sink.send_error_async(e)
        except Exception as e:
            rpc_error = slim_bindings.RpcError.Rpc(
                code=slim_bindings.RpcCode.INTERNAL,
                message=str(e),
                details=None
            )
            await sink.send_error_async(rpc_error)"#;

const STREAM_UNARY_HANDLER_CLASS_TEMPLATE: &str = r#"

class _{{SERVICE_NAME}}Servicer_{{METHOD_NAME}}_Handler:
    def __init__(self, servicer):
        self.servicer = servicer

    async def handle(self, stream: slim_bindings.RequestStream, context: slim_bindings.Context) -> bytes:
        try:
            async def request_iterator():
                while True:
                    stream_msg = await stream.next_async()
                    if stream_msg.is_end():
                        break
                    if stream_msg.is_error():
                        raise stream_msg[0]
                    if stream_msg.is_data():
                        yield {{INPUT_TYPE_FULL_PATH}}.FromString(stream_msg[0])

            response = await self.servicer.{{METHOD_NAME}}(request_iterator(), context)
            return {{OUTPUT_TYPE_FULL_PATH}}.SerializeToString(response)
        except slim_bindings.RpcError:
            raise
        except Exception as e:
            raise slim_bindings.RpcError.Rpc(
                code=slim_bindings.RpcCode.INTERNAL,
                message=str(e),
                details=None
            )"#;

const STREAM_STREAM_HANDLER_CLASS_TEMPLATE: &str = r#"

class _{{SERVICE_NAME}}Servicer_{{METHOD_NAME}}_Handler:
    def __init__(self, servicer):
        self.servicer = servicer

    async def handle(self, stream: slim_bindings.RequestStream, context: slim_bindings.Context, sink: slim_bindings.ResponseSink):
        try:
            async def request_iterator():
                while True:
                    stream_msg = await stream.next_async()
                    if stream_msg.is_end():
                        break
                    if stream_msg.is_error():
                        raise stream_msg[0]
                    if stream_msg.is_data():
                        yield {{INPUT_TYPE_FULL_PATH}}.FromString(stream_msg[0])

            response_iter = self.servicer.{{METHOD_NAME}}(request_iterator(), context)
            async for response in response_iter:
                await sink.send_async({{OUTPUT_TYPE_FULL_PATH}}.SerializeToString(response))
            await sink.close_async()
        except slim_bindings.RpcError as e:
            await sink.send_error_async(e)
        except Exception as e:
            rpc_error = slim_bindings.RpcError.Rpc(
                code=slim_bindings.RpcCode.INTERNAL,
                message=str(e),
                details=None
            )
            await sink.send_error_async(rpc_error)"#;

const REGISTER_METHOD_TEMPLATE: &str = r#"    server.register_{{RPC_TYPE}}(
        service_name="{{PACKAGE_NAME}}.{{SERVICE_NAME}}",
        method_name="{{METHOD_NAME}}",
        handler=_{{SERVICE_NAME}}Servicer_{{METHOD_NAME}}_Handler(servicer),
    )
"#;

// --- END TEMPLATE DEFINITIONS ---

/// Generate Python slimrpc code from a CodeGeneratorRequest
pub fn generate(request: CodeGeneratorRequest) -> Result<CodeGeneratorResponse> {
    let mut types_module_import = String::new();

    // Parse parameters, if any
    if let Some(p) = &request.parameter {
        // Split by , and process parameters in the form key=value
        for param in p.split(',') {
            let mut parts = param.splitn(2, '=');
            if let (Some(key), Some(value)) = (parts.next(), parts.next()) {
                // Handle parameter
                match key {
                    TYPES_IMPORT => {
                        types_module_import = value.to_string();
                    }
                    _ => {
                        return Err(anyhow::anyhow!("Unknown parameter: {}", key));
                    }
                }
            } else {
                return Err(anyhow::anyhow!("Invalid parameter format: {}", param));
            }
        }
    }

    let mut response = CodeGeneratorResponse::default();

    // Map proto file names to their package names for easy lookup
    let file_to_package: HashMap<String, String> = request
        .proto_file
        .iter()
        .filter_map(|f| {
            f.name
                .as_ref()
                .map(|name| (name.clone(), f.package.clone().unwrap_or_default()))
        })
        .collect();

    // Iterate through each .proto file provided in the request
    for file_descriptor in request.proto_file {
        // Only process files that are explicitly requested for generation
        if !request
            .file_to_generate
            .contains(&file_descriptor.name.clone().unwrap_or_default())
        {
            continue;
        }

        let package_name = file_descriptor.package.clone().unwrap_or_default();
        let file_name_base = file_descriptor
            .name
            .as_ref()
            .and_then(|name| name.strip_suffix(".proto"))
            .map(|name| name.to_snake_case())
            .unwrap_or_else(|| "unknown_proto".to_string());

        // The generated Python file will be named after the proto file,
        // with a '_pb2_slimrpc' suffix
        let output_file_name = format!("{}_pb2_slimrpc.py", file_name_base);

        let mut pb2_imports: HashSet<String> = HashSet::new();
        let mut service_definitions_content = String::new();
        let mut handler_classes_content = String::new();
        let mut add_servicer_functions_content = String::new();
        let mut services_found = false;

        if types_module_import.is_empty() {
            let local_pb2_module = format!("{}_pb2", file_name_base);
            types_module_import = format!("from . import {} as pb2", local_pb2_module);
        }

        pb2_imports.insert(types_module_import.clone());

        // Helper function to resolve the full Python path for a message type
        let mut resolve_type_path = |full_type_name: &str, current_file_package: &str| -> String {
            let trimmed_name = full_type_name.trim_start_matches('.');
            let parts: Vec<&str> = trimmed_name.split('.').collect();
            let type_name = parts.last().unwrap_or(&"").to_string();
            let package_path = parts[..parts.len() - 1].join(".");

            if trimmed_name.starts_with("google.protobuf.") {
                let module_name = format!("google.protobuf.{}_pb2", type_name.to_snake_case());
                let alias_name = module_name.replace('.', "__");
                pb2_imports.insert(format!("import {} as {}", module_name, alias_name));
                return format!("{}.{}", alias_name, type_name);
            }

            if package_path == current_file_package {
                let mut prefix = "pb2.".to_string();

                if !types_module_import.is_empty() {
                    // Get last word from the import and add a .
                    prefix = format!(
                        "{}.",
                        types_module_import
                            .split_whitespace()
                            .last()
                            .unwrap_or(&prefix)
                    );
                }

                format!("{}{}", prefix, type_name)
            } else {
                let mut dep_file_base_name = String::new();
                for (file_name, file_package) in file_to_package.iter() {
                    if file_package == &package_path {
                        dep_file_base_name = file_name
                            .strip_suffix(".proto")
                            .unwrap_or(file_name)
                            .to_string()
                            .to_snake_case();
                        break;
                    }
                }

                if !dep_file_base_name.is_empty() {
                    let dep_pb2_module = format!("{}_pb2", dep_file_base_name);
                    pb2_imports.insert(format!(
                        "from . import {} as {}",
                        dep_pb2_module, dep_pb2_module
                    ));
                    format!("{}.{}", dep_pb2_module, type_name)
                } else {
                    // Fallback to local import if package is not found
                    format!("pb2.{}", type_name)
                }
            }
        };

        // Iterate through services defined in the current .proto file
        for service in file_descriptor.service {
            services_found = true;
            let service_name = service.name.context("Service name missing")?;

            let mut method_stub_methods_content = String::new();
            let mut method_servicers_content = String::new();
            let mut handler_classes_for_service = String::new();
            let mut register_methods_content = String::new();

            // Generate methods for the client stub and server servicer
            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")?;
                let output_type = method.output_type.clone().context("Output type missing")?;

                let input_type_full_path = resolve_type_path(&input_type, &package_name);
                let output_type_full_path = resolve_type_path(&output_type, &package_name);

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

                let rpc_type = match (is_client_streaming, is_server_streaming) {
                    (false, false) => "unary_unary",
                    (false, true) => "unary_stream",
                    (true, false) => "stream_unary",
                    (true, true) => "stream_stream",
                };

                let method_template = if is_client_streaming {
                    STREAM_METHOD_SERVICER_TEMPLATE
                } else {
                    UNARY_METHOD_SERVICER_TEMPLATE
                };

                let handler_class_template = match (is_client_streaming, is_server_streaming) {
                    (false, false) => UNARY_UNARY_HANDLER_CLASS_TEMPLATE,
                    (false, true) => UNARY_STREAM_HANDLER_CLASS_TEMPLATE,
                    (true, false) => STREAM_UNARY_HANDLER_CLASS_TEMPLATE,
                    (true, true) => STREAM_STREAM_HANDLER_CLASS_TEMPLATE,
                };

                let stub_method_template = match (is_client_streaming, is_server_streaming) {
                    (false, false) => UNARY_UNARY_STUB_METHOD_TEMPLATE,
                    (false, true) => UNARY_STREAM_STUB_METHOD_TEMPLATE,
                    (true, false) => STREAM_UNARY_STUB_METHOD_TEMPLATE,
                    (true, true) => STREAM_STREAM_STUB_METHOD_TEMPLATE,
                };

                // Populate method stub method template
                let current_method_stub_method = stub_method_template
                    .replace("{{METHOD_NAME}}", &method_name)
                    .replace("{{PACKAGE_NAME}}", &package_name)
                    .replace("{{SERVICE_NAME}}", &service_name)
                    .replace("{{INPUT_TYPE_FULL_PATH}}", &input_type_full_path)
                    .replace("{{OUTPUT_TYPE_FULL_PATH}}", &output_type_full_path);
                method_stub_methods_content.push_str(&current_method_stub_method);

                // Populate method servicer template
                let current_method_servicer = method_template
                    .replace("{{METHOD_NAME}}", &method_name)
                    .replace("{{INPUT_TYPE_FULL_PATH}}", &input_type_full_path);
                method_servicers_content.push_str(&current_method_servicer);

                // Populate handler class template
                let current_handler_class = handler_class_template
                    .replace("{{SERVICE_NAME}}", &service_name)
                    .replace("{{METHOD_NAME}}", &method_name)
                    .replace("{{INPUT_TYPE_FULL_PATH}}", &input_type_full_path)
                    .replace("{{OUTPUT_TYPE_FULL_PATH}}", &output_type_full_path);
                handler_classes_for_service.push_str(&current_handler_class);

                // Populate register method template
                let current_register_method = REGISTER_METHOD_TEMPLATE
                    .replace("{{SERVICE_NAME}}", &service_name)
                    .replace("{{METHOD_NAME}}", &method_name)
                    .replace("{{RPC_TYPE}}", rpc_type)
                    .replace("{{PACKAGE_NAME}}", &package_name);
                register_methods_content.push_str(&current_register_method);
            }

            // Populate service stub template
            let current_service_stub = SERVICE_STUB_TEMPLATE
                .replace("{{SERVICE_NAME}}", &service_name)
                .replace("{{METHOD_STUB_METHODS}}", &method_stub_methods_content);
            service_definitions_content.push_str(&current_service_stub);

            // Populate service servicer template
            let current_service_servicer = SERVICE_SERVICER_TEMPLATE
                .replace("{{SERVICE_NAME}}", &service_name)
                .replace("{{METHOD_SERVICERS}}", &method_servicers_content);
            service_definitions_content.push_str(&current_service_servicer);

            // Add handler classes for this service
            handler_classes_content.push_str(&handler_classes_for_service);

            // Populate add servicer function template
            let current_add_servicer_function = ADD_SERVICER_FUNCTION_TEMPLATE
                .replace("{{SERVICE_NAME}}", &service_name)
                .replace("{{REGISTER_METHODS}}", register_methods_content.trim_end());
            add_servicer_functions_content.push_str(&current_add_servicer_function);
        }

        // Only add a file to the response if services were found and generated
        if services_found {
            let pb2_imports_string = pb2_imports.into_iter().collect::<Vec<_>>().join("\n");
            let final_content = FILE_TEMPLATE
                .replace("{{PB2_IMPORTS}}", &pb2_imports_string)
                .replace("{{SERVICE_DEFINITIONS}}", &service_definitions_content)
                .replace("{{HANDLER_CLASSES}}", &handler_classes_content)
                .replace(
                    "{{ADD_SERVICER_FUNCTIONS}}",
                    &add_servicer_functions_content,
                );

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

    Ok(response)
}

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

    /// Helper function to create a basic FileDescriptorProto for testing
    fn create_test_file_descriptor(
        name: &str,
        package: &str,
        services: Vec<ServiceDescriptorProto>,
    ) -> FileDescriptorProto {
        FileDescriptorProto {
            name: Some(name.to_string()),
            package: Some(package.to_string()),
            service: services,
            message_type: vec![],
            enum_type: vec![],
            dependency: vec![],
            public_dependency: vec![],
            weak_dependency: vec![],
            extension: vec![],
            options: None,
            source_code_info: None,
            syntax: None,
        }
    }

    /// Helper function to create a basic MethodDescriptorProto for testing
    fn create_test_method(
        name: &str,
        input_type: &str,
        output_type: &str,
        client_streaming: bool,
        server_streaming: bool,
    ) -> MethodDescriptorProto {
        MethodDescriptorProto {
            name: Some(name.to_string()),
            input_type: Some(input_type.to_string()),
            output_type: Some(output_type.to_string()),
            client_streaming: Some(client_streaming),
            server_streaming: Some(server_streaming),
            options: None,
        }
    }

    /// Helper function to create a basic ServiceDescriptorProto for testing
    fn create_test_service(
        name: &str,
        methods: Vec<MethodDescriptorProto>,
    ) -> ServiceDescriptorProto {
        ServiceDescriptorProto {
            name: Some(name.to_string()),
            method: methods,
            options: None,
        }
    }

    #[test]
    fn test_generate_no_services() {
        // Test with a proto file that has no services
        let request = CodeGeneratorRequest {
            file_to_generate: vec!["test.proto".to_string()],
            parameter: None,
            proto_file: vec![create_test_file_descriptor(
                "test.proto",
                "test.package",
                vec![],
            )],
            compiler_version: None,
        };

        let response = generate(request).unwrap();

        // Should not generate any files since there are no services
        assert_eq!(response.file.len(), 0);
    }

    #[test]
    fn test_generate_with_unary_unary_method() {
        // Test with a simple unary-unary RPC method
        let method = create_test_method(
            "GetUser",
            ".test.package.GetUserRequest",
            ".test.package.GetUserResponse",
            false,
            false,
        );
        let service = create_test_service("UserService", vec![method]);
        let file_descriptor =
            create_test_file_descriptor("user.proto", "test.package", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["user.proto".to_string()],
            parameter: None,
            proto_file: vec![file_descriptor],
            compiler_version: None,
        };

        let response = generate(request).unwrap();

        assert_eq!(response.file.len(), 1);
        let generated_file = &response.file[0];
        assert_eq!(generated_file.name.as_ref().unwrap(), "user_pb2_slimrpc.py");

        let content = generated_file.content.as_ref().unwrap();
        assert!(content.contains("class UserServiceStub:"));
        assert!(content.contains("class UserServiceServicer:"));
        assert!(content.contains("def GetUser(self, request, context):"));
        assert!(content.contains(
            "def add_UserServiceServicer_to_server(servicer, server: slim_bindings.Server):"
        ));
        assert!(content.contains("unary_unary"));
        // Verify metadata parameter is included in stub method signature with type annotations
        assert!(content.contains("async def GetUser(self, request: pb2.GetUserRequest, timeout: Optional[timedelta] = None, metadata: Optional[dict[str, str]] = None)"));
        // Verify imports are present
        assert!(content.contains("from datetime import timedelta"));
        assert!(content.contains("from typing import Optional"));
    }

    #[test]
    fn test_generate_with_streaming_methods() {
        // Test with all streaming combinations
        let methods = vec![
            create_test_method(
                "UnaryToStream",
                ".pkg.Request",
                ".pkg.Response",
                false,
                true,
            ),
            create_test_method(
                "StreamToUnary",
                ".pkg.Request",
                ".pkg.Response",
                true,
                false,
            ),
            create_test_method(
                "StreamToStream",
                ".pkg.Request",
                ".pkg.Response",
                true,
                true,
            ),
        ];
        let service = create_test_service("StreamService", methods);
        let file_descriptor = create_test_file_descriptor("stream.proto", "pkg", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["stream.proto".to_string()],
            parameter: None,
            proto_file: vec![file_descriptor],
            compiler_version: None,
        };

        let response = generate(request).unwrap();

        assert_eq!(response.file.len(), 1);
        let content = response.file[0].content.as_ref().unwrap();

        assert!(content.contains("unary_stream"));
        assert!(content.contains("stream_unary"));
        assert!(content.contains("stream_stream"));
        assert!(content.contains("def UnaryToStream(self, request, context):"));
        assert!(content.contains("def StreamToUnary(self, request_iterator, context):"));
        assert!(content.contains("def StreamToStream(self, request_iterator, context):"));

        // Verify metadata parameter is included in all stub method signatures with type annotations
        assert!(content.contains("async def UnaryToStream(self, request: pb2.Request, timeout: Optional[timedelta] = None, metadata: Optional[dict[str, str]] = None):"));
        assert!(content.contains("async def StreamToUnary(self, request_iterator, timeout: Optional[timedelta] = None, metadata: Optional[dict[str, str]] = None)"));
        assert!(content.contains("async def StreamToStream(self, request_iterator, timeout: Optional[timedelta] = None, metadata: Optional[dict[str, str]] = None):"));
        // Verify imports are present
        assert!(content.contains("from datetime import timedelta"));
        assert!(content.contains("from typing import Optional"));
    }

    #[test]
    fn test_generate_with_types_import_parameter() {
        // Test with custom types_import parameter
        let method = create_test_method("Echo", ".test.Echo", ".test.Echo", false, false);
        let service = create_test_service("TestService", vec![method]);
        let file_descriptor = create_test_file_descriptor("test.proto", "test", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["test.proto".to_string()],
            parameter: Some("types_import=from mymodule import types as pb2".to_string()),
            proto_file: vec![file_descriptor],
            compiler_version: None,
        };

        let response = generate(request).unwrap();

        assert_eq!(response.file.len(), 1);
        let content = response.file[0].content.as_ref().unwrap();
        assert!(content.contains("from mymodule import types as pb2"));
    }

    #[test]
    fn test_generate_with_invalid_parameter() {
        // Test with an unknown parameter
        let method = create_test_method("Test", ".test.Req", ".test.Res", false, false);
        let service = create_test_service("TestService", vec![method]);
        let file_descriptor = create_test_file_descriptor("test.proto", "test", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["test.proto".to_string()],
            parameter: Some("unknown_param=value".to_string()),
            proto_file: vec![file_descriptor],
            compiler_version: None,
        };

        let result = generate(request);
        assert!(result.is_err());
        assert!(
            result
                .unwrap_err()
                .to_string()
                .contains("Unknown parameter")
        );
    }

    #[test]
    fn test_generate_with_invalid_parameter_format() {
        // Test with invalid parameter format (no equals sign)
        let method = create_test_method("Test", ".test.Req", ".test.Res", false, false);
        let service = create_test_service("TestService", vec![method]);
        let file_descriptor = create_test_file_descriptor("test.proto", "test", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["test.proto".to_string()],
            parameter: Some("invalid_format".to_string()),
            proto_file: vec![file_descriptor],
            compiler_version: None,
        };

        let result = generate(request);
        assert!(result.is_err());
        assert!(
            result
                .unwrap_err()
                .to_string()
                .contains("Invalid parameter format")
        );
    }

    #[test]
    fn test_generate_multiple_services() {
        // Test with multiple services in the same file
        let method1 = create_test_method("Method1", ".pkg.Req1", ".pkg.Res1", false, false);
        let method2 = create_test_method("Method2", ".pkg.Req2", ".pkg.Res2", false, false);

        let service1 = create_test_service("Service1", vec![method1]);
        let service2 = create_test_service("Service2", vec![method2]);

        let file_descriptor =
            create_test_file_descriptor("multi.proto", "pkg", vec![service1, service2]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["multi.proto".to_string()],
            parameter: None,
            proto_file: vec![file_descriptor],
            compiler_version: None,
        };

        let response = generate(request).unwrap();

        assert_eq!(response.file.len(), 1);
        let content = response.file[0].content.as_ref().unwrap();

        assert!(content.contains("class Service1Stub:"));
        assert!(content.contains("class Service1Servicer:"));
        assert!(content.contains("class Service2Stub:"));
        assert!(content.contains("class Service2Servicer:"));
        assert!(content.contains("def add_Service1Servicer_to_server"));
        assert!(content.contains("def add_Service2Servicer_to_server"));
    }

    #[test]
    fn test_generate_not_in_files_to_generate() {
        // Test that files not in file_to_generate are skipped
        let method = create_test_method("Test", ".pkg.Req", ".pkg.Res", false, false);
        let service = create_test_service("TestService", vec![method]);
        let file_descriptor = create_test_file_descriptor("test.proto", "pkg", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["other.proto".to_string()], // Different file
            parameter: None,
            proto_file: vec![file_descriptor],
            compiler_version: None,
        };

        let response = generate(request).unwrap();

        // Should not generate any files since test.proto is not in file_to_generate
        assert_eq!(response.file.len(), 0);
    }

    #[test]
    fn test_generate_snake_case_file_name() {
        // Test that file names are converted to snake_case
        let method = create_test_method("Test", ".pkg.Req", ".pkg.Res", false, false);
        let service = create_test_service("TestService", vec![method]);
        let file_descriptor =
            create_test_file_descriptor("CamelCaseFile.proto", "pkg", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["CamelCaseFile.proto".to_string()],
            parameter: None,
            proto_file: vec![file_descriptor],
            compiler_version: None,
        };

        let response = generate(request).unwrap();

        assert_eq!(response.file.len(), 1);
        assert_eq!(
            response.file[0].name.as_ref().unwrap(),
            "camel_case_file_pb2_slimrpc.py"
        );
    }

    #[test]
    fn test_generate_empty_package_name() {
        // Test with empty package name
        let method = create_test_method("Test", ".Req", ".Res", false, false);
        let service = create_test_service("TestService", vec![method]);
        let file_descriptor = create_test_file_descriptor("test.proto", "", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["test.proto".to_string()],
            parameter: None,
            proto_file: vec![file_descriptor],
            compiler_version: None,
        };

        let response = generate(request).unwrap();

        assert_eq!(response.file.len(), 1);
        // Should still generate valid code
        let content = response.file[0].content.as_ref().unwrap();
        assert!(content.contains("class TestServiceStub:"));
    }

    #[test]
    fn test_generate_multiple_parameters() {
        // Test with multiple comma-separated parameters
        let method = create_test_method("Test", ".test.Req", ".test.Res", false, false);
        let service = create_test_service("TestService", vec![method]);
        let file_descriptor = create_test_file_descriptor("test.proto", "test", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["test.proto".to_string()],
            parameter: Some("types_import=from custom import pb2".to_string()),
            proto_file: vec![file_descriptor],
            compiler_version: None,
        };

        let response = generate(request).unwrap();

        assert_eq!(response.file.len(), 1);
        let content = response.file[0].content.as_ref().unwrap();
        assert!(content.contains("from custom import pb2"));
    }

    #[test]
    fn test_generate_google_protobuf_types() {
        // Test with google.protobuf types (like Empty)
        let method = create_test_method(
            "TestEmpty",
            ".google.protobuf.Empty",
            ".test.Response",
            false,
            false,
        );
        let service = create_test_service("TestService", vec![method]);
        let file_descriptor = create_test_file_descriptor("test.proto", "test", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["test.proto".to_string()],
            parameter: None,
            proto_file: vec![file_descriptor],
            compiler_version: None,
        };

        let response = generate(request).unwrap();

        assert_eq!(response.file.len(), 1);
        let content = response.file[0].content.as_ref().unwrap();
        // Should import google.protobuf types correctly
        assert!(content.contains("google.protobuf.empty_pb2"));
    }

    #[test]
    fn test_generate_service_with_multiple_methods() {
        // Test service with multiple methods
        let methods = vec![
            create_test_method("Create", ".pkg.CreateReq", ".pkg.CreateRes", false, false),
            create_test_method("Read", ".pkg.ReadReq", ".pkg.ReadRes", false, false),
            create_test_method("Update", ".pkg.UpdateReq", ".pkg.UpdateRes", false, false),
            create_test_method("Delete", ".pkg.DeleteReq", ".pkg.DeleteRes", false, false),
        ];
        let service = create_test_service("CrudService", methods);
        let file_descriptor = create_test_file_descriptor("crud.proto", "pkg", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["crud.proto".to_string()],
            parameter: None,
            proto_file: vec![file_descriptor],
            compiler_version: None,
        };

        let response = generate(request).unwrap();

        assert_eq!(response.file.len(), 1);
        let content = response.file[0].content.as_ref().unwrap();

        assert!(content.contains("async def Create("));
        assert!(content.contains("async def Read("));
        assert!(content.contains("async def Update("));
        assert!(content.contains("async def Delete("));
        assert!(content.contains("call_unary_async"));
    }

    #[test]
    fn test_generate_cross_package_dependency() {
        // Test the else branch where package_path != current_file_package (line 214)
        // This tests types from a different package that IS found in file_to_package

        // Create a dependency proto file with its own package
        let dep_file_descriptor = create_test_file_descriptor(
            "common/types.proto",
            "common.types",
            vec![], // No services in dependency
        );

        // Create the main service that uses types from the dependency package
        let method = create_test_method(
            "ProcessRequest",
            ".common.types.CommonRequest", // Type from different package
            ".service.pkg.ServiceResponse", // Type from current package
            false,
            false,
        );
        let service = create_test_service("MyService", vec![method]);
        let main_file_descriptor =
            create_test_file_descriptor("service/api.proto", "service.pkg", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["service/api.proto".to_string()],
            parameter: None,
            proto_file: vec![dep_file_descriptor, main_file_descriptor],
            compiler_version: None,
        };

        let response = generate(request).unwrap();

        assert_eq!(response.file.len(), 1);
        let content = response.file[0].content.as_ref().unwrap();

        // Should import the dependency package
        assert!(content.contains("from . import common_types_pb2 as common_types_pb2"));
        // Should use the imported module for the cross-package type
        assert!(content.contains("common_types_pb2.CommonRequest"));
        // Should use pb2 prefix for local package type
        assert!(content.contains("pb2.ServiceResponse"));
    }

    #[test]
    fn test_generate_cross_package_type_not_found() {
        // Test the fallback else case when package is not found (lines 234-237)
        // This happens when a type references a package not in file_to_package map
        let method = create_test_method(
            "ProcessData",
            ".unknown.package.Request", // Package not in proto_file list
            ".test.package.Response",
            false,
            false,
        );
        let service = create_test_service("DataService", vec![method]);
        let file_descriptor =
            create_test_file_descriptor("data.proto", "test.package", vec![service]);

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["data.proto".to_string()],
            parameter: None,
            proto_file: vec![file_descriptor],
            compiler_version: None,
        };

        let response = generate(request).unwrap();

        assert_eq!(response.file.len(), 1);
        let content = response.file[0].content.as_ref().unwrap();

        // When package is not found, it should fallback to local pb2 import
        assert!(content.contains("pb2.Request"));
        assert!(content.contains("pb2.Request.FromString"));
        assert!(content.contains("class DataServiceStub:"));
    }
}
