package swift

import (
	"path/filepath"
	"strings"

	"github.com/bazelbuild/bazel-gazelle/config"
	"github.com/bazelbuild/bazel-gazelle/language"
	"github.com/bazelbuild/bazel-gazelle/language/proto"
	"github.com/bazelbuild/bazel-gazelle/rule"
	"golang.org/x/text/cases"

	lang "golang.org/x/text/language"
)

// RulesFromProtos returns the Bazel build rule declarations for the provided source files.
func RulesFromProtos(
	args language.GenerateArgs,
	swiftProtoGenerationMode string,
	generateSwiftProtoLibraries bool,
	generateSwiftProtoLibraryGRPCFlavors []string,
	swiftProtoCompilers map[string]string,
	generatePascalCaseModuleNames bool,
	omitProtoSuffixFromModuleNames bool,
) []*rule.Rule {

	// Extract information about proto_library targets generated by other languages.
	// This is a collection of proto_library rule names that have a corresponding
	// swift_proto_library rule already generated.
	protoPackageNamesToProtoLibraryNames := make(map[string][]string)
	protoLibraryTargetNamesToProtoPackages := make(map[string]proto.Package)
	for _, r := range args.OtherGen {
		if r.Kind() != "proto_library" {
			continue
		}
		protoPackage := r.PrivateAttr(proto.PackageKey).(proto.Package)
		protoPackageNamesToProtoLibraryNames[protoPackage.Name] =
			append(protoPackageNamesToProtoLibraryNames[protoPackage.Name], r.Name())
		protoLibraryTargetNamesToProtoPackages[r.Name()] = protoPackage
	}

	// Generate the rules from proto packages:
	var rules []*rule.Rule
	switch swiftProtoGenerationMode {
	case "match":
		for protoLibraryTargetName, protoPackage := range protoLibraryTargetNamesToProtoPackages {

			// Generate the rule from the proto package:
			swiftProtoLibraryBaseName := strings.TrimSuffix(protoLibraryTargetName, "_proto")
			rs := generateRuleFromProtoPackage(
				args,
				swiftProtoLibraryBaseName,
				[]string{protoLibraryTargetName},
				protoPackage,
				generateSwiftProtoLibraries,
				generateSwiftProtoLibraryGRPCFlavors,
				swiftProtoCompilers,
				generatePascalCaseModuleNames,
				omitProtoSuffixFromModuleNames,
			)
			rules = append(rules, rs...)
		}
	case "package":
		for protoPackageName, protoLibraryTargetNames := range protoPackageNamesToProtoLibraryNames {
			// Merge the proto packages:
			mergedProtoPackage := proto.Package{
				Name:        protoPackageName,
				RuleName:    "",
				Files:       make(map[string]proto.FileInfo),
				Imports:     make(map[string]bool),
				Options:     make(map[string]string),
				HasServices: false,
			}
			for _, protoLibraryTargetName := range protoLibraryTargetNames {
				protoPackage := protoLibraryTargetNamesToProtoPackages[protoLibraryTargetName]
				for file, fileInfo := range protoPackage.Files {
					mergedProtoPackage.Files[file] = fileInfo
				}
				for protoImport, protoImportValue := range protoPackage.Imports {
					mergedProtoPackage.Imports[protoImport] = protoImportValue
				}
				for option, optionValue := range protoPackage.Options {
					mergedProtoPackage.Options[option] = optionValue
				}
				mergedProtoPackage.HasServices = mergedProtoPackage.HasServices || protoPackage.HasServices
			}

			// Generate the rule from the merged package:
			swiftProtoLibraryBaseName := protoPackageName
			rs := generateRuleFromProtoPackage(
				args,
				swiftProtoLibraryBaseName,
				protoLibraryTargetNames,
				mergedProtoPackage,
				generateSwiftProtoLibraries,
				generateSwiftProtoLibraryGRPCFlavors,
				swiftProtoCompilers,
				generatePascalCaseModuleNames,
				omitProtoSuffixFromModuleNames,
			)
			rules = append(rules, rs...)
		}
	default:
		panic("unknown swift proto generation mode: " + swiftProtoGenerationMode)
	}

	return rules
}

func pascalCaseModuleName(snakeCaseModuleName string) string {
	moduleNameComponents := strings.Split(snakeCaseModuleName, "_")
	caser := cases.Title(lang.English)
	pascalCaseModuleName := ""
	for _, component := range moduleNameComponents {
		pascalCaseModuleName += caser.String(component)
	}

	return pascalCaseModuleName
}

func deriveSwiftProtoLibraryTargetNameAndModuleName(
	swiftProtoLibraryBaseName string,
	protoFlavor string,
	generatePascalCaseModuleNames bool,
	omitProtoSuffixFromModuleNames bool,
) (string, string) {

	// Determine the swift_proto_library target name:
	swiftProtoLibraryTargetName := swiftProtoLibraryBaseName
	swiftProtoLibraryTargetName += "_" + protoFlavor

	// Determine the swift_proto_library module name:
	swiftProtoLibraryModuleName := swiftProtoLibraryBaseName + strings.TrimPrefix(protoFlavor, "swift")
	if omitProtoSuffixFromModuleNames {
		swiftProtoLibraryModuleName = strings.TrimSuffix(swiftProtoLibraryModuleName, "_proto")
	}
	if generatePascalCaseModuleNames {
		swiftProtoLibraryModuleName = pascalCaseModuleName(swiftProtoLibraryModuleName)
	}

	return swiftProtoLibraryTargetName, swiftProtoLibraryModuleName
}

// These protos are statically linked into the SwiftProtobuf module:
// https://github.com/apple/swift-protobuf/blob/main/Sources/SwiftProtobufPluginLibrary/SwiftProtobufInfo.swift
var wellKnownTypeProtoImports map[string]struct{} = map[string]struct{}{
	"google/protobuf/timestamp.proto":      {},
	"google/protobuf/field_mask.proto":     {},
	"google/protobuf/api.proto":            {},
	"google/protobuf/duration.proto":       {},
	"google/protobuf/struct.proto":         {},
	"google/protobuf/wrappers.proto":       {},
	"google/protobuf/source_context.proto": {},
	"google/protobuf/any.proto":            {},
	"google/protobuf/type.proto":           {},
	"google/protobuf/empty.proto":          {},
	"google/protobuf/descriptor.proto":     {},
}

func isWellKnownTypeProtoImport(protoImport string) bool {
	_, ok := wellKnownTypeProtoImports[protoImport]
	return ok
}

func generateRuleFromProtoPackage(
	args language.GenerateArgs,
	swiftProtoLibraryBaseName string,
	protoLibraryTargetNames []string,
	protoPackage proto.Package,
	generateSwiftProtoLibraries bool,
	generateSwiftProtoLibraryGRPCFlavors []string,
	swiftProtoCompilers map[string]string,
	generatePascalCaseModuleNames bool,
	omitProtoSuffixFromModuleNames bool,
) []*rule.Rule {

	// Map proto target names to labels:
	protoLibraryTargetLabels := []string{}
	for _, protoLibraryTargetName := range protoLibraryTargetNames {
		protoLibraryTargetLabels = append(protoLibraryTargetLabels, ":"+protoLibraryTargetName)
	}

	// Determine whether we should set visibility on the target:
	shouldSetVisibility := shouldSetVisibility(args)

	// Create the swift proto package:
	swiftProtoPackage := SwiftProtoPackage{
		Dir:          args.Dir,
		Rel:          args.Rel,
		ProtoPackage: protoPackage,
	}

	// Build a map of proto paths provided by this package:
	providedImports := map[string]struct{}{}
	for protoFileName := range swiftProtoPackage.ProtoPackage.Files {
		protoPath := filepath.Join(swiftProtoPackage.Rel, protoFileName)
		providedImports[protoPath] = struct{}{}
	}

	// Build the list of protos imported by the proto package which we need to resolve:
	gazelleImports := []string{}
	for protoImport := range protoPackage.Imports {
		if isWellKnownTypeProtoImport(protoImport) {
			// Skip these protos because they're statically linked into SwiftProtobuf.
			continue
		} else if _, ok := providedImports[protoImport]; ok {
			// Skip these protos because they are provided by the current package.
			continue
		} else {
			// These are just the paths to the specific protos to import, e.g.:
			// "protos/echoservice/messages/echo_messages.proto"
			// We need to defer resolving these to swift_proto_library targets until initial generation is complete.
			gazelleImports = append(gazelleImports, protoImport)
		}
	}

	// Generate the swift_proto_library:
	rules := []*rule.Rule{}
	if protoPackage.HasServices {

		// Generate the targets for the requested variants:
		for _, protoFlavor := range generateSwiftProtoLibraryGRPCFlavors {

			// Determine the target name and module name for the swift_proto_library target:
			swiftProtoLibraryTargetName, swiftProtoLibraryModuleName := deriveSwiftProtoLibraryTargetNameAndModuleName(
				swiftProtoLibraryBaseName,
				protoFlavor,
				generatePascalCaseModuleNames,
				omitProtoSuffixFromModuleNames,
			)

			// Generate the swift_proto_library target for the flavor:
			swiftProtoLibrary := rule.NewRule("swift_proto_library", swiftProtoLibraryTargetName)
			swiftProtoLibrary.SetAttr(ModuleNameAttrName, swiftProtoLibraryModuleName)
			swiftProtoLibrary.SetAttr("protos", protoLibraryTargetLabels)
			swiftProtoLibrary.SetAttr("compilers", []string{
				swiftProtoCompilers["swift_proto"],
				swiftProtoCompilers[protoFlavor],
			})
			swiftProtoLibrary.SetPrivateAttr(config.GazelleImportsKey, gazelleImports)
			swiftProtoLibrary.SetPrivateAttr(SwiftProtoPackageKey, swiftProtoPackage)
			setVisibilityAttr(swiftProtoLibrary, shouldSetVisibility, []string{"//visibility:public"})
			rules = append(rules, swiftProtoLibrary)
		}
	} else if generateSwiftProtoLibraries {

		// Determine the target name and module name for the swift_proto_library target:
		swiftProtoLibraryTargetName, swiftProtoLibraryModuleName := deriveSwiftProtoLibraryTargetNameAndModuleName(
			swiftProtoLibraryBaseName,
			"swift_proto",
			generatePascalCaseModuleNames,
			omitProtoSuffixFromModuleNames,
		)

		// Generate the standard, non-GRPC swift_proto_library target:
		swiftProtoLibrary := rule.NewRule("swift_proto_library", swiftProtoLibraryTargetName)
		swiftProtoLibrary.SetAttr(ModuleNameAttrName, swiftProtoLibraryModuleName)
		swiftProtoLibrary.SetAttr("protos", protoLibraryTargetLabels)
		swiftProtoLibrary.SetAttr("compilers", []string{swiftProtoCompilers["swift_proto"]})
		swiftProtoLibrary.SetPrivateAttr(config.GazelleImportsKey, gazelleImports)
		swiftProtoLibrary.SetPrivateAttr(SwiftProtoPackageKey, swiftProtoPackage)
		setVisibilityAttr(swiftProtoLibrary, shouldSetVisibility, []string{"//visibility:public"})
		rules = append(rules, swiftProtoLibrary)
	}

	return rules
}
