// Package dnsresolve resolves every subdomain present in the store, plus
// any that arrive via late SubdomainDiscovered events while the module is
// running. Results (IPs, CNAME, PTR) are written back to the store AND
// announced via DNSResolved events for downstream enrichment modules.
//
// This module is idempotent: Upsert on the same subdomain twice is cheap.
package dnsresolve

import (
	"context"
	"strings"
	"sync"
	"time"

	"god-eye/internal/config"
	godns "god-eye/internal/dns"
	"god-eye/internal/eventbus"
	"god-eye/internal/module"
	"god-eye/internal/store"
)

const ModuleName = "dns.resolver"

type resolverModule struct{}

func Register() { module.Register(&resolverModule{}) }

func (*resolverModule) Name() string                   { return ModuleName }
func (*resolverModule) Phase() module.Phase            { return module.PhaseResolution }
func (*resolverModule) Consumes() []eventbus.EventType { return []eventbus.EventType{eventbus.EventSubdomainDiscovered} }
func (*resolverModule) Produces() []eventbus.EventType { return []eventbus.EventType{eventbus.EventDNSResolved} }
func (*resolverModule) DefaultEnabled() bool           { return true }

func (m *resolverModule) Run(mctx module.Context) error {
	resolvers := parseResolvers(mctx.Config.String("resolvers", ""))
	timeout := mctx.Config.Int("timeout", 5)
	conc := mctx.Config.Int("concurrency", 500)
	if conc <= 0 {
		conc = 500
	}

	// Dedup across drain + late events.
	processed := make(map[string]struct{})
	var processedMu sync.Mutex
	shouldProcess := func(sub string) bool {
		processedMu.Lock()
		defer processedMu.Unlock()
		if _, dup := processed[sub]; dup {
			return false
		}
		processed[sub] = struct{}{}
		return true
	}

	work := make(chan string, conc*2)
	var wg sync.WaitGroup
	for i := 0; i < conc; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for sub := range work {
				m.resolveOne(mctx, sub, resolvers, timeout)
			}
		}()
	}

	// 1) Drain the store: every subdomain discovered so far goes in.
	for _, h := range mctx.Store.All(mctx.Ctx) {
		if h == nil || h.Subdomain == "" {
			continue
		}
		if !shouldProcess(h.Subdomain) {
			continue
		}
		select {
		case work <- h.Subdomain:
		case <-mctx.Ctx.Done():
			close(work)
			wg.Wait()
			return nil
		}
	}

	// 2) Keep listening for late events (e.g. from recursive discovery that
	// runs in our own phase and produces new subdomains mid-resolution).
	sub := mctx.Bus.Subscribe(eventbus.EventSubdomainDiscovered, func(_ context.Context, e eventbus.Event) {
		ev, ok := e.(eventbus.SubdomainDiscovered)
		if !ok {
			return
		}
		if !shouldProcess(ev.Subdomain) {
			return
		}
		select {
		case work <- ev.Subdomain:
		case <-mctx.Ctx.Done():
		}
	})
	defer sub.Unsubscribe()

	// 3) Give late events a short window to arrive (e.g. recursive module
	// running concurrently in PhaseResolution). 1 second is enough — we
	// already drained the store, so any straggler events here are rare.
	select {
	case <-time.After(1 * time.Second):
	case <-mctx.Ctx.Done():
	}

	close(work)
	wg.Wait()
	return nil
}

func (m *resolverModule) resolveOne(mctx module.Context, sub string, resolvers []string, timeout int) {
	if err := mctx.Ctx.Err(); err != nil {
		return
	}
	ips := godns.ResolveSubdomain(sub, resolvers, timeout)
	if len(ips) == 0 {
		return
	}

	cname := godns.ResolveCNAME(sub, resolvers, timeout)
	ptr := godns.ResolvePTR(ips[0], resolvers, timeout)

	_ = mctx.Store.Upsert(mctx.Ctx, sub, func(h *store.Host) {
		store.AddIPs(h, ips)
		if cname != "" && h.CNAME == "" {
			h.CNAME = cname
		}
		if ptr != "" && h.PTR == "" {
			h.PTR = ptr
		}
		store.AddDiscoveryMethod(h, "resolved")
	})

	mctx.Bus.Publish(mctx.Ctx, eventbus.DNSResolved{
		EventMeta: eventbus.EventMeta{At: time.Now(), Source: ModuleName, Target: sub},
		Subdomain: sub,
		IPs:       ips,
		CNAME:     cname,
		PTR:       ptr,
	})
}

func parseResolvers(s string) []string {
	s = strings.TrimSpace(s)
	if s == "" {
		return config.DefaultResolvers
	}
	var out []string
	for _, r := range strings.Split(s, ",") {
		r = strings.TrimSpace(r)
		if r == "" {
			continue
		}
		if !strings.Contains(r, ":") {
			r = r + ":53"
		}
		out = append(out, r)
	}
	if len(out) == 0 {
		return config.DefaultResolvers
	}
	return out
}