package ingesters

import (
	"context"
	"encoding/json"

	"github.com/deepfence/ThreatMapper/deepfence_utils/directory"
	"github.com/deepfence/ThreatMapper/deepfence_utils/log"
	"github.com/deepfence/ThreatMapper/deepfence_utils/utils"
	ingestersUtil "github.com/deepfence/ThreatMapper/deepfence_utils/utils/ingesters"
	"github.com/twmb/franz-go/pkg/kgo"
)

type VulnerabilityIngester struct{}

func NewVulnerabilityIngester() KafkaIngester[[]ingestersUtil.Vulnerability] {
	return &VulnerabilityIngester{}
}

func (tc *VulnerabilityIngester) Ingest(
	ctx context.Context,
	cs []ingestersUtil.Vulnerability,
	ingestC chan *kgo.Record,
) error {

	tenantID, err := directory.ExtractNamespace(ctx)
	if err != nil {
		return err
	}

	rh := []kgo.RecordHeader{
		{Key: "namespace", Value: []byte(tenantID)},
	}

	for _, c := range cs {
		cb, err := json.Marshal(c)
		if err != nil {
			log.Error().Msg(err.Error())
		} else {
			ingestC <- &kgo.Record{
				Topic:   utils.TopicWithNamespace(utils.VulnerabilityScan, string(tenantID)),
				Value:   cb,
				Headers: rh,
			}
		}
	}

	return nil
}

type VulnerabilityStatusIngester struct{}

func NewVulnerabilityStatusIngester() KafkaIngester[[]ingestersUtil.VulnerabilityScanStatus] {
	return &VulnerabilityStatusIngester{}
}

func (tc *VulnerabilityStatusIngester) Ingest(
	ctx context.Context,
	statuses []ingestersUtil.VulnerabilityScanStatus,
	ingestC chan *kgo.Record,
) error {
	tenantID, err := directory.ExtractNamespace(ctx)
	if err != nil {
		return err
	}

	rh := []kgo.RecordHeader{
		{Key: "namespace", Value: []byte(tenantID)},
	}

	for _, c := range statuses {
		cb, err := json.Marshal(c)
		if err != nil {
			log.Error().Msg(err.Error())
		} else {
			ingestC <- &kgo.Record{
				Topic:   utils.TopicWithNamespace(utils.VulnerabilityScanStatus, string(tenantID)),
				Value:   cb,
				Headers: rh,
			}
		}
	}

	return nil
}
