package ingesters

import (
	"context"
	"fmt"
	"hash/fnv"
	"time"

	"github.com/deepfence/ThreatMapper/deepfence_utils/directory"
	"github.com/deepfence/ThreatMapper/deepfence_utils/telemetry"
	"github.com/deepfence/ThreatMapper/deepfence_utils/utils"
	ingestersUtil "github.com/deepfence/ThreatMapper/deepfence_utils/utils/ingesters"
	"github.com/neo4j/neo4j-go-driver/v5/neo4j"
)

func generateMalwareRuleId(metaRule ingestersUtil.MetaRules) string {
	return fmt.Sprintf("malware-%s", metaRule.RuleName)
}

func generateHashFromString(s string) string {
	h := fnv.New32a()
	h.Write([]byte(s))
	return fmt.Sprintf("%x", h.Sum32())
}

func CommitFuncMalware(ctx context.Context, ns string, data []ingestersUtil.Malware) error {
	ctx = directory.ContextWithNameSpace(ctx, directory.NamespaceID(ns))

	ctx, span := telemetry.NewSpan(ctx, "ingesters", "commit-func-malware")
	defer span.End()

	driver, err := directory.Neo4jClient(ctx)
	if err != nil {
		return err
	}

	session := driver.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeWrite})
	defer session.Close(ctx)

	tx, err := session.BeginTransaction(ctx, neo4j.WithTxTimeout(30*time.Second))
	if err != nil {
		return err
	}
	defer tx.Close(ctx)

	dataMap, err := malwareToMaps(data)
	if err != nil {
		return nil
	}

	if _, err = tx.Run(ctx, `
		UNWIND $batch as row WITH row.Rule as rule, row.Malware as malware
		MATCH (r:MalwareRule{rule_id:rule.rule_id})
		WITH malware as row, r
		MERGE (n:Malware{node_id:row.node_id})
		SET n+= row,
		    n.masked = COALESCE(n.masked, r.masked, false),
		    n.updated_at = TIMESTAMP()
		WITH n, r, row
		MERGE (n)-[:IS]->(r)
		MERGE (m:MalwareScan{node_id: row.scan_id})
		WITH n, m
		MERGE (m) -[l:DETECTED]-> (n)
		SET l.masked = COALESCE(n.masked, false)`,
		map[string]interface{}{"batch": dataMap}); err != nil {
		return err
	}

	return tx.Commit(ctx)
}

func malwareToMaps(data []ingestersUtil.Malware) ([]map[string]map[string]interface{}, error) {
	var malwares []map[string]map[string]interface{}
	for _, i := range data {
		malware := utils.ToMap(i)
		delete(malware, "meta_rules")

		i.MetaRules.FileSeverity = malware["file_severity"].(string)
		i.MetaRules.RuleName = malware["rule_name"].(string)
		i.MetaRules.RuleID = generateMalwareRuleId(i.MetaRules)

		malware["node_id"] = generateHashFromString(fmt.Sprintf("%v:%v",
			i.MetaRules.RuleID, i.CompleteFilename))
		malware["reference"] = i.MetaRules.Reference

		malwares = append(malwares, map[string]map[string]interface{}{
			"Rule":    utils.ToMap(i.MetaRules),
			"Malware": malware,
		})
	}
	return malwares, nil
}
