package swift

import (
	"bytes"
	"encoding/json"
	"fmt"
	"sort"

	"github.com/bazelbuild/bazel-gazelle/rule"
	mapset "github.com/deckarep/golang-set/v2"
	"golang.org/x/exp/slices"
)

// A DependencyIndex encapsulates the indexes used to lookup Swift modules and products.
type DependencyIndex struct {
	moduleIndex         ModuleIndex
	productIndex        ProductIndex
	modToPrdsIndex      ModuleToProductsIndex
	prdToModsIndex      ProductMembershipsIndex
	directDepIdentities mapset.Set[string]
	packageIndex        PackageIndex
}

// NewDependencyIndex creates an empty dependency index.
func NewDependencyIndex() *DependencyIndex {
	di := &DependencyIndex{
		moduleIndex:         make(ModuleIndex),
		productIndex:        make(ProductIndex),
		directDepIdentities: mapset.NewSet[string](),
		packageIndex:        NewPackageIndex(),
	}
	di.init()
	return di
}

func (di *DependencyIndex) init() {
	di.modToPrdsIndex = make(ModuleToProductsIndex)
	di.prdToModsIndex = make(ProductMembershipsIndex)
	for _, m := range di.moduleIndex.Modules() {
		di.modToPrdsIndex.IndexModule(m)
		di.prdToModsIndex.IndexModule(m)
	}
}

// IndexRepoRule indexes any repository rules that are not already included in the module index
// generated by update-repos. In other words, the swift_package repo rules are not indexed here
// because their contents are already included in the module index JSON file.
func (di *DependencyIndex) IndexRepoRule(r *rule.Rule, repoRoot string) error {
	var err error
	switch r.Kind() {
	case HTTPArchiveRuleKind:
		err = di.indexHTTPArchive(r, repoRoot)
	}
	return err
}

func (di *DependencyIndex) indexHTTPArchive(r *rule.Rule, repoRoot string) error {
	ha, err := NewHTTPArchiveFromRule(r, repoRoot)
	if err != nil {
		return err
	}
	if ha == nil {
		return nil
	}
	di.moduleIndex.Add(ha.Modules...)
	return nil
}

func (di DependencyIndex) AddDirectDependency(identities ...string) {
	for _, id := range identities {
		di.directDepIdentities.Add(id)
	}
}

func (di DependencyIndex) DirectDepIdentities() []string {
	result := di.directDepIdentities.ToSlice()
	slices.Sort(result)
	return result
}

// AddModule adds one or more modules to the underlying indexes.
func (di *DependencyIndex) AddModule(modules ...*Module) {
	for _, m := range modules {
		di.moduleIndex.Add(m)
		di.modToPrdsIndex.IndexModule(m)
		di.prdToModsIndex.IndexModule(m)
	}
}

// AddProduct adds one or more products to the underlying indexes.
func (di *DependencyIndex) AddProduct(products ...*Product) {
	for _, p := range products {
		di.productIndex.Add(p)
	}
}

// IndexBazelRepo indexes the modules and products for a Bazel repository.
func (di *DependencyIndex) IndexBazelRepo(bzlRepo *BazelRepo) error {
	var err error
	pi := bzlRepo.PkgInfo
	repoName := bzlRepo.Name

	// Index modules
	modules := make([]*Module, len(pi.Targets))
	for idx, t := range pi.Targets {
		modules[idx], err = NewModuleFromTarget(repoName, bzlRepo.Identity, t)
		if err != nil {
			return err
		}
	}
	di.AddModule(modules...)

	// Index products
	for _, p := range pi.Products {
		prd, err := NewProductFromPkgInfoProduct(bzlRepo, p)
		if err != nil {
			return err
		}
		di.AddProduct(prd)
	}

	return nil
}

func (di *DependencyIndex) AddPackage(packages ...*Package) {
	di.packageIndex.Add(packages...)
}

func (di *DependencyIndex) GetPackage(identity string) *Package {
	if pkg, ok := di.packageIndex[identity]; ok {
		return pkg
	}
	return nil
}

func (di *DependencyIndex) DirectDepPackages() []*Package {
	directDepIdentities := di.DirectDepIdentities()
	results := make([]*Package, len(directDepIdentities))
	for idx, identity := range directDepIdentities {
		results[idx] = di.GetPackage(identity)
	}
	return results
}

func (di *DependencyIndex) Packages() []*Package {
	return di.packageIndex.Packages()
}

// Resolve

type ModuleResolutionResult struct {
	Products   Products
	Unresolved []string
}

type productSetResult struct {
	pik          ProductIndexKey
	modules      mapset.Set[string]
	intersect    mapset.Set[string]
	intersectCnt int
	extraCnt     int
}

// Implemented to help with debugging resolution issues.
func (psr productSetResult) String() string {
	return fmt.Sprintf(
		"productSetResult{pik: %s, intersectCnt: %d, extraCnt: %d, modules: %s, intersect: %s}",
		psr.pik, psr.intersectCnt, psr.extraCnt, psr.modules, psr.intersect)
}

type productSetResults []productSetResult

func (psrs productSetResults) Len() int {
	return len(psrs)
}

func (psrs productSetResults) Swap(i, j int) {
	psrs[i], psrs[j] = psrs[j], psrs[i]
}

func (psrs productSetResults) Less(i, j int) bool {
	a, b := psrs[i], psrs[j]
	if a.intersectCnt != b.intersectCnt {
		return a.intersectCnt < b.intersectCnt
	}
	// A better match is one where there are fewer extra modules.
	// Hence, those with higher number of extra modules will be less.
	if a.extraCnt != b.extraCnt {
		return a.extraCnt > b.extraCnt
	}
	// Finally, we sort on the product index key. This exists to provide a stable/consistent sort.
	return a.pik < b.pik
}

func (di *DependencyIndex) products(piks ...ProductIndexKey) Products {
	// It is possible to encounter a ProductIndexKey for which there is no product. This is due to
	// phantom products (i.e., inferred by SPM). These phantom products appear in the description
	// JSON, but do not appear in the dump JSON. We opt not to include these phantom products,
	// because they were never specified in the original manifest.
	prds := make(Products, 0, len(piks))
	for _, pik := range piks {
		if p, ok := di.productIndex[pik]; ok {
			prds = append(prds, p)
		}
	}
	return prds
}

// ResolveModulesToProducts finds the best set of products to satisfy the list of modules restricted
// to a set of packages.
func (di *DependencyIndex) ResolveModulesToProducts(
	moduleNames []string,
	pkgIdentities []string,
) *ModuleResolutionResult {

	var result ModuleResolutionResult
	pkgIdentsSet := mapset.NewSet[string](pkgIdentities...)
	unresolved := mapset.NewSet[string]()

	populateUnresolved := func() {
		if unresolved.Cardinality() == 0 {
			result.Unresolved = nil
			return
		}
		result.Unresolved = unresolved.ToSlice()
		sort.Strings(result.Unresolved)
	}

	// Find all of the products that contain a match for any of the modules
	potentialPikSet := mapset.NewSet[ProductIndexKey]()
	for _, mname := range moduleNames {
		if piks, ok := di.modToPrdsIndex[mname]; ok {
			for _, pik := range piks {
				if pkgIdentsSet.Contains(pik.Identity()) {
					potentialPikSet.Add(pik)
				}
			}
		} else {
			unresolved.Add(mname)
		}
	}

	// If we found nothing or if we found only one prodcut, we are done.
	if potentialPikSet.Cardinality() <= 1 {
		result.Products = di.products(potentialPikSet.ToSlice()...)
		populateUnresolved()
		return &result
	}

	var selectedPiks ProductIndexKeys
	populateProducts := func() {
		sort.Sort(selectedPiks)
		result.Products = di.products(selectedPiks...)
	}

	modulesToResolve := mapset.NewSet[string](moduleNames...)
	for modulesToResolve.Cardinality() > 0 {
		pikSetCnt := potentialPikSet.Cardinality()
		if pikSetCnt == 0 {
			// The only reason that we should get here is if we encounter a built-in module name
			// that is not in our list of built-in modules.
			result.Unresolved = append(result.Unresolved, modulesToResolve.ToSlice()...)
			populateProducts()
			populateUnresolved()
			return &result
		}
		psrs := make(productSetResults, 0, pikSetCnt)
		for pik := range potentialPikSet.Iterator().C {
			var psr productSetResult
			psr.pik = pik
			psr.modules = di.prdToModsIndex[pik]
			psr.intersect = psr.modules.Intersect(modulesToResolve)
			psr.intersectCnt = psr.intersect.Cardinality()
			psr.extraCnt = psr.modules.Cardinality() - psr.intersectCnt
			psrs = append(psrs, psr)
		}
		sort.Sort(sort.Reverse(psrs))
		selectedPsr := psrs[0]

		if selectedPsr.intersectCnt > 0 {
			// We found a product that contains at least one of the modules that we are trying to
			// resolve. Add it to the list of selected products and remove the modules that it
			// contains from the list of modules that we are trying to resolve.
			selectedPiks = append(selectedPiks, selectedPsr.pik)
			modulesToResolve = modulesToResolve.Difference(selectedPsr.intersect)
		}

		potentialPikSet.Remove(selectedPsr.pik)
	}

	populateProducts()
	populateUnresolved()
	return &result
}

func (di *DependencyIndex) FindModules(moduleName string, pkgIdentities []string) Modules {
	compFn := func(a, b *Module) int {
		if a.PkgIdentity < b.PkgIdentity {
			return -1
		}
		if a.PkgIdentity == b.PkgIdentity {
			return 0
		}
		return 1
	}
	modules, ok := di.moduleIndex[moduleName]
	if !ok {
		return nil
	}
	if len(pkgIdentities) == 0 {
		// We are going to sort the results. Create a copy so that we do not affect the underlying
		// slice.
		results := make(Modules, len(modules))
		copy(results, modules)
		slices.SortFunc(results, compFn)
		return results
	}
	piSet := mapset.NewSet[string](pkgIdentities...)
	results := make(Modules, 0, len(modules))
	for _, m := range modules {
		if piSet.Contains(m.PkgIdentity) {
			results = append(results, m)
		}
	}
	slices.SortFunc(results, compFn)
	return results
}

// JSON

type dependencyIndexJSONData struct {
	DirectDepIdentities []string   `json:"direct_dep_identities"`
	Modules             []*Module  `json:"modules"`
	Products            []*Product `json:"products"`
	Packages            []*Package `json:"packages,omitempty"`
}

func (di *DependencyIndex) jsonData() *dependencyIndexJSONData {
	return &dependencyIndexJSONData{
		DirectDepIdentities: di.DirectDepIdentities(),
		Modules:             di.moduleIndex.Modules(),
		Products:            di.productIndex.Products(),
		Packages:            di.packageIndex.Packages(),
	}
}

func newFromJSONData(jd *dependencyIndexJSONData) *DependencyIndex {
	di := &DependencyIndex{
		directDepIdentities: mapset.NewSet[string](jd.DirectDepIdentities...),
		moduleIndex:         NewModuleIndex(jd.Modules...),
		productIndex:        NewProductIndex(jd.Products...),
		packageIndex:        NewPackageIndex(jd.Packages...),
	}
	di.init()
	return di
}

func (di *DependencyIndex) MarshalJSON() ([]byte, error) {
	jd := di.jsonData()
	return json.Marshal(jd)
}

func (di *DependencyIndex) UnmarshalJSON(b []byte) error {
	var jd dependencyIndexJSONData
	if err := json.Unmarshal(b, &jd); err != nil {
		return err
	}
	newdi := newFromJSONData(&jd)
	newdi.init()
	*di = *newdi
	return nil
}

// NewDependencyIndexFromJSON creates a dependency index from JSON.
func NewDependencyIndexFromJSON(data []byte) (*DependencyIndex, error) {
	var di DependencyIndex
	if err := json.Unmarshal(data, &di); err != nil {
		return nil, err
	}
	return &di, nil
}

// JSON returns the pretty, JSON representation for the dependency index that is written to disk.
func (di *DependencyIndex) JSON() ([]byte, error) {
	b, err := json.Marshal(di)
	if err != nil {
		return nil, err
	}
	var buf bytes.Buffer
	err = json.Indent(&buf, b, "", "  ")
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}
