package main

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/bazelbuild/buildtools/build"

	container "edge-infra.dev/hack/build/rules/container/gazelle/language"
	"edge-infra.dev/pkg/lib/build/bazel"
	"edge-infra.dev/pkg/lib/cli/sink"
)

func New() *sink.Command {
	return &sink.Command{
		Use:   "tp-migrate [space separated list of files to migrate]",
		Short: "Migrate .bzl files containing dictionaries to use third_party_container_dep macros",
		Exec: func(_ context.Context, r sink.Run) error {
			filesToMigrate := r.Args()
			wd := bazel.ResolveWdOrDie()

			for _, fPath := range filesToMigrate {
				if !strings.Contains(fPath, ".bzl") {
					r.Log.Error(errors.New("skipping non-bzl file"), "filePath", fPath)
					continue
				}
				r.Log.Info("found", "bzl file", fPath)
				migratedPath, migrationRequired, err := migrateBzlFile(wd, fPath, r)
				if err != nil {
					return err
				}

				if migrationRequired {
					r.Log.Info("migrated", "file", migratedPath)
				} else {
					r.Log.Info("no image dictionaries present to migrate", "file", fPath)
				}
			}

			return nil
		},
	}
}

// Main function body to migrate a .bzl file's image dictionaries to a wrapper function
// containing third_party_container_dep rules instead
func migrateBzlFile(wdPath string, fPath string, r sink.Run) (string, bool, error) {
	migratedFileName := fmt.Sprintf("%s_migrated%s", strings.TrimSuffix(filepath.Base(fPath), filepath.Ext(fPath)), filepath.Ext(fPath))
	migratedFilePath := fmt.Sprintf("%s/%s/%s", wdPath, filepath.Dir(fPath), migratedFileName)

	originalPath := filepath.Join(wdPath, fPath)
	bzlBytes, err := os.ReadFile(originalPath)
	if err != nil {
		r.Log.Error(err, "error reading original file", "file path", originalPath)
		return "", false, err
	}

	bzlFile, err := build.ParseBzl(originalPath, bzlBytes)
	if err != nil {
		r.Log.Error(err, "error parsing bzl file", "bzl file path", originalPath)
		return "", false, err
	}

	imageDictReplaceMap := map[int]build.Expr{}
	for i, s := range bzlFile.Stmt {
		if assignExpr, isAssign := s.(*build.AssignExpr); isAssign {
			if dictStmt, isDict := assignExpr.RHS.(*build.DictExpr); isDict {
				wrapperDefName := strings.ReplaceAll(assignExpr.LHS.(*build.Ident).Name, "-", "_")
				wrapperDefName = strings.ToLower(wrapperDefName)

				wrapperDef := imageDictToTpDepWrapperDef(wrapperDefName, dictStmt)

				imageDictReplaceMap[i] = wrapperDef
			}
		}
	}
	if len(imageDictReplaceMap) == 0 {
		return "", false, nil
	}

	migratedFile := bzlFile.Copy().(*build.File)

	for replaceIndex, wrapperDef := range imageDictReplaceMap {
		migratedFile.Stmt[replaceIndex] = wrapperDef
	}

	tpDepLoad := &build.LoadStmt{
		Module:       &build.StringExpr{Value: "//hack/build/rules/container:third_party_images.bzl"},
		From:         []*build.Ident{&build.Ident{Name: container.ThirdPartyContainerDepRuleName}},
		To:           []*build.Ident{&build.Ident{Name: container.ThirdPartyContainerDepRuleName}},
		ForceCompact: true,
	}
	migratedFile.Stmt = append(migratedFile.Stmt, tpDepLoad)
	err = os.WriteFile(migratedFilePath, build.Format(migratedFile), 0644)

	build.Format(migratedFile)
	if err != nil {
		r.Log.Error(err, "error writing file", "file", migratedFilePath)
		return "", false, err
	}

	return migratedFilePath, true, nil
}

// imageDictToTpDepWrappeDef returns a *build.DefExpr that represents a wrapper function
// containing all third_party_container_dep calls previously represented as dictionary
// items
func imageDictToTpDepWrapperDef(wrapperDefName string, dict *build.DictExpr) *build.DefStmt {
	defExpr := &build.DefStmt{Name: wrapperDefName}

	for _, baseKVExpr := range dict.List {
		callExpr := &build.CallExpr{
			X:              &build.Ident{Name: container.ThirdPartyContainerDepRuleName},
			ForceMultiLine: true,
		}

		ruleName := strings.ReplaceAll(baseKVExpr.Key.(*build.StringExpr).Value, "-", "_")

		assignExprs := dictValsToAssignExprs(ruleName, baseKVExpr.Value.(*build.DictExpr))
		callExpr.List = append(callExpr.List, assignExprs...)

		defExpr.Body = append(defExpr.Body, callExpr)
	}
	return defExpr
}

// Convert dictionary values to AssignExprs
func dictValsToAssignExprs(name string, dictExpr *build.DictExpr) []build.Expr {
	assignExprs := []build.Expr{
		&build.AssignExpr{
			LHS: &build.Ident{Name: "name"},
			Op:  "=",
			RHS: &build.StringExpr{Value: name},
		},
	}

	for _, e := range dictExpr.List {
		attr := strings.ToLower(e.Key.(*build.StringExpr).Value)
		attr = attrRenames(attr)
		assignExpr := &build.AssignExpr{
			LHS: &build.Ident{Name: attr},
			Op:  "=",
			RHS: e.Value,
		}

		assignExprs = append(assignExprs, assignExpr)
	}
	return assignExprs
}

// Any attributes that might need to be renamed go here
func attrRenames(attr string) string {
	renameMap := map[string]string{
		"repo": "repository",
	}

	if rename, found := renameMap[attr]; found {
		return rename
	}
	return attr
}