package engine

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/tidbcloud/infra-provider/infra-inspector/api"
	"github.com/tidbcloud/infra-provider/infra-inspector/pkg/config"
	"github.com/tidbcloud/infra-provider/infra-inspector/pkg/inspect/lister"

	"github.com/open-policy-agent/opa/rego"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
)

const (
	queryRegoTemplate = "output := data.%s.violation"
)

var (
	_ Engine = &regoEngine{}
)

type regoEngine struct {
	cfg   config.RegoEngineConfig
	rego  *rego.Rego
	query rego.PreparedEvalQuery
}

func NewRegoEngine(cfg config.RegoEngineConfig) (*regoEngine, error) {
	rego, err := initRego(cfg.RegoRulePath)
	if err != nil {
		return nil, fmt.Errorf("init rego failed: %v", err)
	}

	query, err := rego.PrepareForEval(context.TODO())
	if err != nil {
		return nil, fmt.Errorf("failed to prepare query: %v", err)
	}

	r := &regoEngine{
		cfg:   cfg,
		rego:  rego,
		query: query,
	}

	return r, nil
}

func (e *regoEngine) Validate(ctx context.Context, _ lister.ResourceLister, resource *unstructured.Unstructured) (*api.Violation, error) {
	regoResults, err := e.query.Eval(ctx, rego.EvalInput(resource.Object))
	if err != nil {
		return nil, fmt.Errorf("failed to evaluate query: %v", err)
	}

	if len(regoResults) == 0 {
		return nil, fmt.Errorf("no results found")
	}

	jsonResult, err := json.Marshal(regoResults[0].Bindings["output"])
	if err != nil {
		return nil, fmt.Errorf("failed to marshal raw result: %v", err)
	}

	rules := []api.ViolatedRule{}
	err = json.Unmarshal(jsonResult, &rules)
	if err != nil {
		return nil, fmt.Errorf("failed to unmarshal to result: %v", err)
	}

	result := &api.Violation{
		Resource:      GetResourceFromUnstructured(resource),
		ViolatedRules: rules,
	}

	return result, nil
}

func initRego(path string) (*rego.Rego, error) {
	fi, err := os.Stat(path)
	if err != nil {
		return nil, err
	}

	pkgName := ""
	opts := []func(*rego.Rego){}
	if fi.IsDir() {
		entrys, err := os.ReadDir(path)
		if err != nil {
			return nil, err
		}

		for _, entry := range entrys {
			if !entry.Type().IsRegular() {
				continue
			}

			if filepath.Ext(entry.Name()) == ".rego" {
				pkg, m, err := readRuleFromFile(filepath.Join(path, entry.Name()))
				if err != nil {
					return nil, err
				}
				if pkgName != "" && pkgName != pkg {
					return nil, fmt.Errorf("found multiple packages in the same rule dir")
				}
				pkgName = pkg
				opts = append(opts, m)
			}
		}
	} else {
		pkg, m, err := readRuleFromFile(path)
		if err != nil {
			return nil, err
		}
		pkgName = pkg
		opts = append(opts, m)
	}

	if pkgName == "" {
		return nil, fmt.Errorf("package name is empty")
	}
	if len(opts) == 0 {
		return nil, fmt.Errorf("no rego rule found")
	}

	opts = append(opts, rego.Query(fmt.Sprintf(queryRegoTemplate, pkgName)))
	rego := rego.New(opts...)
	return rego, nil
}

func readRuleFromFile(path string) (string /*package*/, func(*rego.Rego), error) {
	pkg, err := readPackageNameFromRule(path)
	if err != nil {
		return "", nil, err
	}

	// read rule
	rule, err := os.ReadFile(path)
	if err != nil {
		return "", nil, err
	}

	return pkg, rego.Module(path, string(rule)), nil
}

func readPackageNameFromRule(path string) (string, error) {
	file, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)

	if scanner.Scan() {
		fields := strings.Fields(scanner.Text())
		if fields[0] != "package" || len(fields) < 2 {
			return "", fmt.Errorf("package name not found in the first line")
		}
		return fields[1], nil
	}

	if scanner.Err() != nil {
		return "", err
	}

	return "", fmt.Errorf("package name not found in the first line")
}