package coredns

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/coredns/caddy/caddyfile"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

// updateCorefile updates the CoreDNS corefile contents.
type UpdateCoreDNS interface {
	UpdateCoreDNS(ctx context.Context, cli client.Client, dnsMap map[string]string) error
	CoreDNSData(ctx context.Context, cli client.Client) (string, error)
}

// hostRecord is a pair of string denoting ip to hostname mapping.
type hostRecord struct {
	ip   string
	host string
}

const (
	hostsKey           = "hosts"
	fallthroughKey     = "fallthrough"
	fallthroughLineStr = "[" + fallthroughKey + "]"
	corednsName        = "coredns"
	corefileKey        = "Corefile"
)

// parseCorefile is an internal function for corefile parsing
func parseCorefile(cf string) (caddyfile.EncodedCaddyfile, error) {
	res := caddyfile.EncodedCaddyfile{}
	buf, err := caddyfile.ToJSON([]byte(cf))
	if err != nil {
		return nil, err
	}
	err = json.Unmarshal(buf, &res)
	if err != nil {
		return nil, err
	}
	return res, nil
}

// insertHostRecord is an internal function for adding DNS record to hostfile in corefile
func insertHostRecord(cf caddyfile.EncodedCaddyfile, records []hostRecord) error {
	newRecords, hostset := make([]interface{}, len(records)), make(map[string]bool)
	for i, record := range records {
		newRecords[i] = []interface{}{record.ip, record.host}
		hostset[record.host] = true
	}
	for idx, server := range cf {
		hostBlkFound, hostIndex, needFallthrough := false, -1, false
		for i, block := range server.Body {
			if len(block) == 0 || fmt.Sprintf("%v", block[0]) != hostsKey {
				continue
			}
			hostBlkFound, hostIndex = true, i
			break
		}
		var hostBlk, entries, newEntries []interface{}
		if !hostBlkFound {
			hostBlk = []interface{}{hostsKey}
			cf[idx].Body = append(server.Body, hostBlk)
			server = cf[idx]
			hostIndex = len(server.Body) - 1
		} else {
			hostBlk = server.Body[hostIndex]
		}
		// create hostfile lists if not created yet
		if len(hostBlk) < 2 {
			entries = []interface{}{}
			server.Body[hostIndex] = append(hostBlk, entries)
			hostBlk = server.Body[hostIndex]
		} else {
			ents, ok := hostBlk[1].([]interface{})
			if !ok {
				return fmt.Errorf("failed to cast hostfile entries to list")
			}
			entries = ents
		}
		if len(entries) == 0 || fmt.Sprintf("%v", entries[len(entries)-1]) != fallthroughLineStr {
			needFallthrough = true
		}
		// add the mapping to the top of the host file
		newEntries = append([]interface{}{}, newRecords...)
		for _, entry := range entries {
			// Deduplicate records
			if record, ok := entry.([]interface{}); ok && len(record) == 2 {
				if host, ok := record[1].(string); ok && hostset[host] {
					continue
				}
			}
			newEntries = append(newEntries, entry)
		}
		if needFallthrough {
			newEntries = append(newEntries, []interface{}{fallthroughKey})
		}
		hostBlk[1] = newEntries
	}
	return nil
}

// getCorefileString is an internal function for generating corefile string content
func getCorefileString(cf caddyfile.EncodedCaddyfile) (string, error) {
	buf, err := json.Marshal(&cf)
	if err != nil {
		return "", err
	}
	res, err := caddyfile.FromJSON(buf)
	if err != nil {
		return "", err
	}
	return string(res), nil
}

// addRecordToCorefile is a wrapper for internal logics that adds host record to corefile.
func addRecordToCorefile(corefile string, records []hostRecord) (string, error) {
	cf, err := parseCorefile(corefile)
	if err != nil {
		return "", err
	}
	err = insertHostRecord(cf, records)
	if err != nil {
		return "", err
	}
	return getCorefileString(cf)
}

// getHostRecord generate the hostRecord from the dnsMap.
func getHostRecord(dnsMap map[string]string) []hostRecord {
	res := []hostRecord{}
	for host, ip := range dnsMap {
		res = append(res, hostRecord{ip: ip, host: host})
	}
	return res
}

type corednsCM struct{}

// retrieveCorefile returns the CoreDNS corefile contents in string
func retrieveCorefile(ctx context.Context, cli client.Client) (string, *corev1.ConfigMap, error) {
	key := types.NamespacedName{
		Namespace: metav1.NamespaceSystem,
		Name:      corednsName,
	}
	cm := &corev1.ConfigMap{}
	err := cli.Get(ctx, key, cm)
	if err != nil {
		return "", nil, err
	}
	if cm.Data == nil {
		cm.Data = make(map[string]string)
	}
	corefile, exist := cm.Data[corefileKey]
	if !exist {
		return "", nil, fmt.Errorf("corefile content missing for coredns configmap")
	}
	return corefile, cm, nil
}

// UpdateCoreDNS updates the CoreDNS configmap based on dnsMap.
func (c *corednsCM) UpdateCoreDNS(ctx context.Context, cli client.Client, dnsMap map[string]string) error {
	corefile, cm, err := retrieveCorefile(ctx, cli)
	if err != nil {
		return err
	}
	corefile, err = addRecordToCorefile(corefile, getHostRecord(dnsMap))
	if err != nil {
		return err
	}
	cm.Data[corefileKey] = corefile
	return cli.Update(ctx, cm)
}

func (c *corednsCM) CoreDNSData(ctx context.Context, cli client.Client) (string, error) {
	cm := &corev1.ConfigMap{}
	err := cli.Get(ctx, types.NamespacedName{
		Namespace: metav1.NamespaceSystem,
		Name:      corednsName,
	}, cm)
	if err != nil {
		return "", err
	}
	data, err := json.Marshal(cm.Data)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

// NewCoreDNSCM returns the procedure for coredns configmap.
func NewCoreDNSCM() *corednsCM {
	return &corednsCM{}
}

var _ UpdateCoreDNS = &corednsCM{}