// 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;

use crate::common;

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

// --- 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 ---

/// Parse a `go_package` option value into an (import_path, package_alias) pair.
///
/// Handles two forms:
/// - `"github.com/org/repo/pkg"` → alias derived from last path component (`"pkg"`)
/// - `"github.com/org/repo/pkg;alias"` → explicit alias after the semicolon
fn parse_go_package(go_pkg: &str) -> (String, String) {
    if let Some(semi) = go_pkg.find(';') {
        let import_path = go_pkg[..semi].to_string();
        let alias = go_pkg[semi + 1..].to_string();
        return (import_path, alias);
    }
    let alias = go_pkg.split('/').next_back().unwrap_or("pb").to_string();
    (go_pkg.to_string(), alias)
}

/// Resolve a fully-qualified protobuf type (e.g. `.google.protobuf.Empty`) into:
/// - the Go expression used in generated code (e.g. `emptypb.Empty`)
/// - an optional `(alias, formatted_import_line)` for packages that need a new import
///
/// Types that belong to `current_proto_pkg` are handled via `types_alias` (the
/// `types_import` parameter); all other types are looked up in `type_to_go`.
fn resolve_proto_type(
    qualified_type: &str,
    current_proto_pkg: &str,
    types_alias: &Option<String>,
    type_to_go: &HashMap<String, (String, String)>,
) -> (String, Option<(String, String)>) {
    let trimmed = qualified_type.trim_start_matches('.');
    let bare = trimmed
        .split('.')
        .next_back()
        .unwrap_or(trimmed)
        .to_string();

    // The proto package is everything before the last dot-component.
    let type_proto_pkg = trimmed.rfind('.').map(|pos| &trimmed[..pos]).unwrap_or("");

    if type_proto_pkg == current_proto_pkg || type_proto_pkg.is_empty() {
        // Same package as the file being compiled — apply types_alias if set.
        let qualified = match types_alias {
            Some(alias) => format!("{}.{}", alias, bare),
            None => bare,
        };
        (qualified, None)
    } else if let Some((import_path, go_alias)) = type_to_go.get(trimmed) {
        // External type — add its Go import and qualify the type name.
        // Look up by fully-qualified type name so that types from the same proto
        // package but different Go packages (e.g. google.protobuf.Empty vs
        // google.protobuf.FileDescriptorProto) resolve correctly.
        let last_component = import_path.split('/').next_back().unwrap_or("");
        let import_line = if go_alias == last_component {
            format!("\"{}\"", import_path)
        } else {
            format!("{} \"{}\"", go_alias, import_path)
        };
        (
            format!("{}.{}", go_alias, bare),
            Some((go_alias.clone(), import_line)),
        )
    } else {
        // No Go package info available — fall back to bare name.
        (bare, None)
    }
}

/// Generate Go slimrpc code from a CodeGeneratorRequest
pub fn generate(request: CodeGeneratorRequest) -> Result<CodeGeneratorResponse> {
    let params = common::parse_parameters(request.parameter.as_deref().unwrap_or(""));

    // Extract types_import (Go import path for proto types), stripping any surrounding quotes
    let types_import_path: Option<String> = params
        .get(TYPES_IMPORT)
        .map(|v| v.trim_matches('"').to_string());

    // Determine package alias: explicit types_alias, or last component of import path.
    // Strip surrounding quotes to handle yaml values like types_alias="pb".
    let types_alias: Option<String> = if let Some(ref path) = types_import_path {
        let derived = path.split('/').next_back().unwrap_or("pb").to_string();
        Some(
            params
                .get(TYPES_ALIAS)
                .map(|v| v.trim_matches('"').to_string())
                .unwrap_or(derived),
        )
    } else {
        None
    };

    // Build map: fully-qualified proto type name → (Go import path, Go package alias).
    // Keyed per-type rather than per-package because proto packages can span multiple
    // Go packages (e.g. google.protobuf.Empty lives in emptypb while
    // google.protobuf.FileDescriptorProto lives in descriptorpb).
    let mut type_to_go: HashMap<String, (String, String)> = HashMap::new();
    for file_desc in &request.proto_file {
        let Some(proto_pkg) = &file_desc.package else {
            continue;
        };
        let Some(go_pkg) = file_desc
            .options
            .as_ref()
            .and_then(|o| o.go_package.as_ref())
        else {
            continue;
        };
        let (import_path, alias) = parse_go_package(go_pkg);
        for msg in &file_desc.message_type {
            if let Some(msg_name) = &msg.name {
                let fqn = format!("{}.{}", proto_pkg, msg_name);
                type_to_go
                    .entry(fqn)
                    .or_insert((import_path.clone(), alias.clone()));
            }
        }
    }

    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()
        };

        // Collect Go imports required by this file beyond the fixed set.
        // Key = Go package alias, value = formatted import line (without leading tab/newline).
        let mut extra_imports: HashMap<String, String> = HashMap::new();

        // Seed with types_import if provided by the user.
        if let (Some(path), Some(alias)) = (&types_import_path, &types_alias) {
            let last_component = path.split('/').next_back().unwrap_or("");
            let line = if alias.as_str() == last_component {
                format!("\"{}\"", path)
            } else {
                format!("{} \"{}\"", alias, path)
            };
            extra_imports.insert(alias.clone(), line);
        }

        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 raw_input = method.input_type.clone().context("Input type missing")?;
                let (input_type, input_import) =
                    resolve_proto_type(&raw_input, &package_name, &types_alias, &type_to_go);
                if let Some((alias, line)) = input_import {
                    extra_imports.insert(alias, line);
                }

                let raw_output = method.output_type.clone().context("Output type missing")?;
                let (output_type, output_import) =
                    resolve_proto_type(&raw_output, &package_name, &types_alias, &type_to_go);
                if let Some((alias, line)) = output_import {
                    extra_imports.insert(alias, line);
                }

                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";

            // Sort imports alphabetically for deterministic output.
            let mut import_lines: Vec<&String> = extra_imports.values().collect();
            import_lines.sort();
            let proto_imports: String = import_lines
                .into_iter()
                .map(|s| format!("\t{}\n", s))
                .collect();

            let content = FILE_TEMPLATE
                .replace("{{PROTO_FILE}}", &proto_file_name)
                .replace("{{PACKAGE}}", &go_package)
                .replace("{{PROTO_IMPORTS}}", &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)")
        );
    }

    #[test]
    fn test_generate_without_types_import() {
        // Regression test: no types_import → bare unqualified type names
        let service = ServiceDescriptorProto {
            name: Some("TestService".to_string()),
            method: vec![MethodDescriptorProto {
                name: Some("Call".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],
            parameter: None,
            ..Default::default()
        };

        let response = generate(request).unwrap();
        let content = response.file[0].content.as_ref().unwrap();

        // No import line for types package
        assert!(!content.contains("\"github.com/org/repo/types\""));
        // Bare type names (no package qualifier)
        assert!(content.contains("*Request"));
        assert!(content.contains("*Response"));
        assert!(!content.contains("*types.Request"));
    }

    #[test]
    fn test_generate_with_types_import() {
        // types_import set → import line added, types prefixed with last path component
        let service = ServiceDescriptorProto {
            name: Some("TestService".to_string()),
            method: vec![MethodDescriptorProto {
                name: Some("Call".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],
            parameter: Some("types_import=github.com/org/repo/types".to_string()),
            ..Default::default()
        };

        let response = generate(request).unwrap();
        let content = response.file[0].content.as_ref().unwrap();

        // Import line present (no explicit alias needed since package name matches dir)
        assert!(content.contains("\"github.com/org/repo/types\""));
        // Types prefixed with derived alias "types"
        assert!(content.contains("*types.Request"));
        assert!(content.contains("*types.Response"));
    }

    #[test]
    fn test_generate_with_quoted_types_alias() {
        // types_alias value wrapped in quotes (e.g. from buf.gen.yaml: types_alias="pb")
        // should produce the same output as the unquoted form
        let service = ServiceDescriptorProto {
            name: Some("TestService".to_string()),
            method: vec![MethodDescriptorProto {
                name: Some("Call".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],
            parameter: Some(
                "types_import=github.com/org/repo/types,types_alias=\"pb\"".to_string(),
            ),
            ..Default::default()
        };

        let response = generate(request).unwrap();
        let content = response.file[0].content.as_ref().unwrap();

        // Should produce aliased import without the extra quotes around the alias
        assert!(content.contains("pb \"github.com/org/repo/types\""));
        assert!(!content.contains("\"pb\""));
        assert!(content.contains("*pb.Request"));
    }

    #[test]
    fn test_generate_with_types_import_and_alias() {
        // types_import + types_alias → aliased import and type prefix
        let service = ServiceDescriptorProto {
            name: Some("TestService".to_string()),
            method: vec![MethodDescriptorProto {
                name: Some("Call".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],
            parameter: Some("types_import=github.com/org/repo/types,types_alias=pb".to_string()),
            ..Default::default()
        };

        let response = generate(request).unwrap();
        let content = response.file[0].content.as_ref().unwrap();

        // Aliased import line: pb "github.com/org/repo/types"
        assert!(content.contains("pb \"github.com/org/repo/types\""));
        // Types prefixed with explicit alias "pb"
        assert!(content.contains("*pb.Request"));
        assert!(content.contains("*pb.Response"));
        // No bare type references
        assert!(!content.contains("*types.Request"));
    }

    /// Build a minimal FileDescriptorProto for an imported proto file whose types
    /// are used as method inputs/outputs in another file.
    ///
    /// `message_names` should list every message type defined in the file so that
    /// the per-type lookup map is populated correctly.
    fn create_import_file_descriptor(
        file_name: &str,
        proto_package: &str,
        go_package: &str,
        message_names: &[&str],
    ) -> FileDescriptorProto {
        use prost_types::{DescriptorProto, FileOptions};
        FileDescriptorProto {
            name: Some(file_name.to_string()),
            package: Some(proto_package.to_string()),
            options: Some(FileOptions {
                go_package: Some(go_package.to_string()),
                ..Default::default()
            }),
            message_type: message_names
                .iter()
                .map(|n| DescriptorProto {
                    name: Some(n.to_string()),
                    ..Default::default()
                })
                .collect(),
            ..Default::default()
        }
    }

    #[test]
    fn test_generate_with_imported_proto_type() {
        // Service method uses google.protobuf.Empty from an imported proto file.
        // The generator should auto-detect the Go package and add the import.
        let service = ServiceDescriptorProto {
            name: Some("TestService".to_string()),
            method: vec![MethodDescriptorProto {
                name: Some("Delete".to_string()),
                input_type: Some(".test.Request".to_string()),
                output_type: Some(".google.protobuf.Empty".to_string()),
                client_streaming: Some(false),
                server_streaming: Some(false),
                ..Default::default()
            }],
            ..Default::default()
        };

        let main_file = create_test_file_descriptor("test.proto", "test", vec![service]);
        let empty_file = create_import_file_descriptor(
            "google/protobuf/empty.proto",
            "google.protobuf",
            "google.golang.org/protobuf/types/known/emptypb",
            &["Empty"],
        );

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

        let response = generate(request).unwrap();
        let content = response.file[0].content.as_ref().unwrap();

        // Import added automatically for the external package
        assert!(
            content.contains("\"google.golang.org/protobuf/types/known/emptypb\""),
            "expected emptypb import, got:\n{}",
            content
        );
        // Output type qualified with the Go alias
        assert!(content.contains("*emptypb.Empty"));
        // Input type from current package stays bare
        assert!(content.contains("*Request"));
        assert!(!content.contains("*test.Request"));
    }

    #[test]
    fn test_generate_with_imported_proto_type_and_types_import() {
        // Combines types_import (for same-package types) with an auto-detected
        // import for a type from a different proto file.
        let service = ServiceDescriptorProto {
            name: Some("TestService".to_string()),
            method: vec![MethodDescriptorProto {
                name: Some("Delete".to_string()),
                input_type: Some(".test.Request".to_string()),
                output_type: Some(".google.protobuf.Empty".to_string()),
                client_streaming: Some(false),
                server_streaming: Some(false),
                ..Default::default()
            }],
            ..Default::default()
        };

        let main_file = create_test_file_descriptor("test.proto", "test", vec![service]);
        let empty_file = create_import_file_descriptor(
            "google/protobuf/empty.proto",
            "google.protobuf",
            "google.golang.org/protobuf/types/known/emptypb",
            &["Empty"],
        );

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["test.proto".to_string()],
            proto_file: vec![main_file, empty_file],
            parameter: Some("types_import=github.com/org/repo/types".to_string()),
            ..Default::default()
        };

        let response = generate(request).unwrap();
        let content = response.file[0].content.as_ref().unwrap();

        // Both imports present
        assert!(content.contains("\"github.com/org/repo/types\""));
        assert!(content.contains("\"google.golang.org/protobuf/types/known/emptypb\""));
        // Same-package type uses types_alias, external type uses go_alias
        assert!(content.contains("*types.Request"));
        assert!(content.contains("*emptypb.Empty"));
    }

    #[test]
    fn test_generate_with_go_package_semicolon_form() {
        // go_package with "path;name" format — alias comes from the explicit name.
        let service = ServiceDescriptorProto {
            name: Some("TestService".to_string()),
            method: vec![MethodDescriptorProto {
                name: Some("Call".to_string()),
                input_type: Some(".ext.Msg".to_string()),
                output_type: Some(".test.Response".to_string()),
                client_streaming: Some(false),
                server_streaming: Some(false),
                ..Default::default()
            }],
            ..Default::default()
        };

        let main_file = create_test_file_descriptor("test.proto", "test", vec![service]);
        let ext_file = create_import_file_descriptor(
            "ext/msg.proto",
            "ext",
            "github.com/org/repo/internal/ext;extpb",
            &["Msg"],
        );

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

        let response = generate(request).unwrap();
        let content = response.file[0].content.as_ref().unwrap();

        // The explicit alias "extpb" must appear in the import, not the dir name "ext"
        assert!(
            content.contains("extpb \"github.com/org/repo/internal/ext\""),
            "expected aliased import, got:\n{}",
            content
        );
        assert!(content.contains("*extpb.Msg"));
    }

    #[test]
    fn test_generate_disambiguates_shared_proto_package() {
        // Regression test: two files share the proto package "google.protobuf" but map
        // to different Go packages.  descriptor.proto appears *before* empty.proto in
        // the proto_file list (as buf orders them), so a package-keyed map would
        // incorrectly resolve google.protobuf.Empty to descriptorpb.
        let service = ServiceDescriptorProto {
            name: Some("TestService".to_string()),
            method: vec![MethodDescriptorProto {
                name: Some("Delete".to_string()),
                input_type: Some(".test.Request".to_string()),
                output_type: Some(".google.protobuf.Empty".to_string()),
                client_streaming: Some(false),
                server_streaming: Some(false),
                ..Default::default()
            }],
            ..Default::default()
        };

        let main_file = create_test_file_descriptor("test.proto", "test", vec![service]);
        // descriptor.proto comes first — must NOT poison the lookup for Empty
        let descriptor_file = create_import_file_descriptor(
            "google/protobuf/descriptor.proto",
            "google.protobuf",
            "google.golang.org/protobuf/types/descriptorpb",
            &["FileDescriptorProto", "DescriptorProto"],
        );
        let empty_file = create_import_file_descriptor(
            "google/protobuf/empty.proto",
            "google.protobuf",
            "google.golang.org/protobuf/types/known/emptypb",
            &["Empty"],
        );

        let request = CodeGeneratorRequest {
            file_to_generate: vec!["test.proto".to_string()],
            // descriptor.proto intentionally first, matching buf's ordering
            proto_file: vec![main_file, descriptor_file, empty_file],
            ..Default::default()
        };

        let response = generate(request).unwrap();
        let content = response.file[0].content.as_ref().unwrap();

        // Must import emptypb, not descriptorpb
        assert!(
            content.contains("\"google.golang.org/protobuf/types/known/emptypb\""),
            "expected emptypb import, got:\n{}",
            content
        );
        assert!(
            !content.contains("descriptorpb"),
            "unexpected descriptorpb import, got:\n{}",
            content
        );
        assert!(content.contains("*emptypb.Empty"));
    }
}
