package threatintel

import (
	"archive/tar"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"strings"

	"github.com/deepfence/ThreatMapper/deepfence_utils/directory"
	"github.com/deepfence/ThreatMapper/deepfence_utils/log"
	"github.com/deepfence/ThreatMapper/deepfence_utils/telemetry"
	"github.com/jellydator/ttlcache/v3"
	"github.com/neo4j/neo4j-go-driver/v5/neo4j"
)

const (
	MalwareRulesStore = "malware"
)

func DownloadMalwareRules(ctx context.Context, entry Entry) error {

	log.Info().Msg("download latest malware rules")

	ctx, span := telemetry.NewSpan(ctx, "threatintel", "download-malware-rules")
	defer span.End()

	// remove old rule file
	existing, _, err := FetchMalwareRulesInfo(ctx)
	if err != nil {
		log.Error().Err(err).Msg("no existing malware rules info found")
	} else {
		if err := DeleteFileMinio(ctx, existing); err != nil {
			log.Error().Err(err).Msgf("failed to delete file %s", existing)
		}
	}

	// download latest rules and uplaod to minio
	content, err := downloadFile(ctx, entry.URL)
	if err != nil {
		log.Error().Err(err).Msg("failed to download malware rules")
		return err
	}

	path, sha, err := UploadToMinio(ctx, content.Bytes(),
		MalwareRulesStore, fmt.Sprintf("malware-rules-%d.tar.gz", entry.Built.Unix()))
	if err != nil {
		log.Error().Err(err).Msg("failed to upload malware rules to fileserver")
		return err
	}

	err = IngestMalwareRules(ctx, content.Bytes())
	if err != nil {
		log.Error().Err(err).Msg("failed to ingest malware rules")
		return err
	}

	// create node in neo4j
	return UpdateMalwareRulesInfo(ctx, sha, strings.TrimPrefix(path, "database/"))
}

func UpdateMalwareRulesInfo(ctx context.Context, hash, path string) error {
	nc, err := directory.Neo4jClient(ctx)
	if err != nil {
		return err
	}
	session := nc.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeWrite})
	defer session.Close(ctx)

	_, err = session.Run(ctx, `
	MERGE (n:MalwareRules{node_id: "latest"})
	SET n.rules_hash=$hash,
		n.path=$path,
		n.updated_at=TIMESTAMP()`,
		map[string]interface{}{
			"hash": hash,
			"path": path,
		})
	if err != nil {
		log.Error().Err(err).Msg("failed to update MalwareRules on neo4j")
		return err
	}

	return nil
}

func FetchMalwareRulesURL(ctx context.Context, consoleURL string, ttlCache *ttlcache.Cache[string, string]) (string, string, error) {
	path, hash, err := FetchMalwareRulesInfo(ctx)
	if err != nil {
		return "", "", err
	}
	exposedURL, err := ExposeFile(ctx, path, consoleURL, ttlCache)
	if err != nil {
		log.Error().Err(err).Msg("failed to expose malware rules on fileserver")
		return "", "", err
	}
	return exposedURL, hash, nil
}

func FetchMalwareRulesInfo(ctx context.Context) (path, hash string, err error) {
	nc, err := directory.Neo4jClient(ctx)
	if err != nil {
		return "", "", err
	}
	session := nc.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeWrite})
	defer session.Close(ctx)

	tx, err := session.BeginTransaction(ctx)
	if err != nil {
		return "", "", err
	}
	defer tx.Close(ctx)

	queryMalwareRules := `
	MATCH (s:MalwareRules{node_id: "latest"})
	RETURN s.path, s.rules_hash`

	r, err := tx.Run(ctx, queryMalwareRules, map[string]interface{}{})
	if err != nil {
		return "", "", err
	}
	rec, err := r.Single(ctx)
	if err != nil {
		return "", "", err
	}

	return rec.Values[0].(string), rec.Values[1].(string), nil
}

func IngestMalwareRules(ctx context.Context, content []byte) error {
	nc, err := directory.Neo4jClient(ctx)
	if err != nil {
		return err
	}
	session := nc.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeWrite})
	defer session.Close(ctx)

	err = ProcessTarGz(content, func(header *tar.Header, reader io.Reader) error {
		var feeds FeedsBundle
		if header.FileInfo().IsDir() {
			return nil
		}
		if strings.HasSuffix(header.Name, ".data") {
			return nil
		}
		jdec := json.NewDecoder(reader)
		err = jdec.Decode(&feeds)
		if err != nil {
			log.Warn().Msg(err.Error())
			return nil
		}

		log.Info().Msgf("Ingesting %d malware", len(feeds.ScannerFeeds.MalwareRules))

		_, err = session.Run(ctx, `
			UNWIND $rules as row
			MERGE (n:DeepfenceRule:MalwareRule{rule_id: row.rule_id})
			SET n.type = row.type,
				n.payload = row.payload,
				n.summary = row.description,
				n.severity = row.severity,
				n.updated_at = TIMESTAMP()`,
			map[string]interface{}{
				"rules": DeepfenceRule2json(feeds.ScannerFeeds.MalwareRules),
			})
		return err
	})
	if err != nil {
		return err
	}

	return err
}
