package resolver

import (
	"encoding/json"
	"fmt"
	"net"
	"os"
	"sync"
	"time"
)

// ServiceResolver maps IP:port:proto to service names
type ServiceResolver struct {
	mu           sync.RWMutex
	portToSvc    map[string]string
	servicesFile string
}

func NewServiceResolver(servicesFile string) *ServiceResolver {
	r := &ServiceResolver{servicesFile: servicesFile, portToSvc: make(map[string]string)}
	r.reload()
	go r.watchReload()
	return r
}

func (r *ServiceResolver) ServiceForDst(ip net.IP, port uint16, proto string) string {
	r.mu.RLock()
	defer r.mu.RUnlock()

	// Try IP:port:proto first
	if svc, ok := r.portToSvc[fmt.Sprintf("%s:%d:%s", ip.String(), port, proto)]; ok {
		return svc
	}
	// Fall back to IP only
	if svc, ok := r.portToSvc[ip.String()]; ok {
		return svc
	}
	return ""
}

func (r *ServiceResolver) reload() {
	data, err := os.ReadFile(r.servicesFile)
	if err != nil {
		return
	}

	var services map[string]interface{}
	if json.Unmarshal(data, &services) != nil {
		return
	}

	newMap := make(map[string]string)
	for name, svcRaw := range services {
		svc, ok := svcRaw.(map[string]interface{})
		if !ok {
			continue
		}

		hosts := map[string]bool{}
		if hostsRaw, ok := svc["hosts"].(map[string]interface{}); ok {
			for ip := range hostsRaw {
				hosts[ip] = true
				newMap[ip] = name
			}
		}

		if portsRaw, ok := svc["ports"].([]interface{}); ok {
			for _, portRaw := range portsRaw {
				port, ok := portRaw.(map[string]interface{})
				if !ok {
					continue
				}
				portNum := fmt.Sprintf("%.0f", port["port"])
				proto, _ := port["proto"].(string)
				for ip := range hosts {
					newMap[fmt.Sprintf("%s:%s:%s", ip, portNum, proto)] = name
				}
			}
		}
	}

	r.mu.Lock()
	r.portToSvc = newMap
	r.mu.Unlock()
}

func (r *ServiceResolver) watchReload() {
	ticker := time.NewTicker(60 * time.Second)
	defer ticker.Stop()
	for range ticker.C {
		r.reload()
	}
}