// Package selfip discovers the host's current WAN address so the SSRF
// guard can deny hairpin-NAT fetches of the residential IP.
package selfip

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"net"
	"net/netip"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"git.wntrmute.dev/kyle/sift/internal/guard"
)

const (
	DefaultInterval     = 10 * time.Minute
	DefaultKeepPrevious = 24 * time.Hour
	DefaultPath         = "/var/lib/sift/self_ip.json"
	DefaultStartupWait  = 10 * time.Second
	DefaultTrace        = "trace:https://1.1.1.1/cdn-cgi/trace"
)

// LookupIP resolves a hostname. Tests inject this.
type LookupIP = guard.LookupIP

// GetFunc GETs a URL (must go through the guarded fetcher) and returns the body.
type GetFunc func(ctx context.Context, rawURL string) ([]byte, error)

// Entry is one observed self address.
type Entry struct {
	Addr   string    `json:"addr"`
	Source string    `json:"source"`
	Seen   time.Time `json:"seen"`
}

type file struct {
	Entries []Entry `json:"entries"`
}

// Status is the /healthz self_ip object.
type Status struct {
	Addrs  []string  `json:"addrs"`
	Age    string    `json:"age"`
	Source string    `json:"source"`
	At     time.Time `json:"at,omitempty"`
}

// Config is the [self_ip] block.
type Config struct {
	Sources      []string
	Interval     time.Duration
	KeepPrevious time.Duration
	Path         string
	StartupWait  time.Duration
}

// Discoverer updates Guard.SetSelfDeny from DDNS and/or a Cloudflare-style trace.
type Discoverer struct {
	Cfg    Config
	Lookup LookupIP
	Get    GetFunc
	Log    *slog.Logger
	Guard  *guard.Guard

	Failures atomic.Uint64
	now      func() time.Time

	mu         sync.Mutex
	entries    []Entry
	lastOK     time.Time
	lastSource string
}

func (d *Discoverer) clock() time.Time {
	if d != nil && d.now != nil {
		return d.now()
	}
	return time.Now()
}

func (d *Discoverer) log() *slog.Logger {
	if d != nil && d.Log != nil {
		return d.Log
	}
	return slog.Default()
}

func (d *Discoverer) interval() time.Duration {
	if d.Cfg.Interval > 0 {
		return d.Cfg.Interval
	}
	return DefaultInterval
}

func (d *Discoverer) keep() time.Duration {
	if d.Cfg.KeepPrevious > 0 {
		return d.Cfg.KeepPrevious
	}
	return DefaultKeepPrevious
}

func (d *Discoverer) path() string {
	if d.Cfg.Path != "" {
		return d.Cfg.Path
	}
	return DefaultPath
}

func (d *Discoverer) sources() []string {
	if len(d.Cfg.Sources) > 0 {
		return d.Cfg.Sources
	}
	return []string{DefaultTrace}
}

// Load reads the persist file. Missing file is not an error.
func (d *Discoverer) Load() error {
	b, err := os.ReadFile(d.path())
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}
	var f file
	if err := json.Unmarshal(b, &f); err != nil {
		return err
	}
	d.mu.Lock()
	d.entries = f.Entries
	if len(d.entries) > 0 {
		d.lastOK = d.entries[len(d.entries)-1].Seen
		d.lastSource = d.entries[len(d.entries)-1].Source
	}
	d.mu.Unlock()
	d.applyLocked()
	return nil
}

// Persist writes the current entries.
func (d *Discoverer) Persist() error {
	d.mu.Lock()
	ents := append([]Entry(nil), d.entries...)
	d.mu.Unlock()
	b, err := json.MarshalIndent(file{Entries: ents}, "", "  ")
	if err != nil {
		return err
	}
	tmp := d.path() + ".tmp"
	if err := os.WriteFile(tmp, b, 0o640); err != nil {
		return err
	}
	return os.Rename(tmp, d.path())
}

// Status is safe for /healthz.
func (d *Discoverer) Status() Status {
	if d == nil {
		return Status{}
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	st := Status{Source: d.lastSource, At: d.lastOK}
	seen := map[string]struct{}{}
	for _, e := range d.entries {
		p := prefixOf(e.Addr)
		if p == "" {
			continue
		}
		if _, ok := seen[p]; ok {
			continue
		}
		seen[p] = struct{}{}
		st.Addrs = append(st.Addrs, p)
	}
	if !d.lastOK.IsZero() {
		st.Age = d.clock().Sub(d.lastOK).Truncate(time.Second).String()
	}
	return st
}

// RunOnce tries each source in order and unions successful global-unicast results.
// On total failure it leaves the existing set in place, counts a failure, and returns the error.
func (d *Discoverer) RunOnce(ctx context.Context) error {
	now := d.clock()
	var got []Entry
	var first string
	var lastErr error
	for _, src := range d.sources() {
		addrs, err := d.resolve(ctx, src)
		if err != nil {
			lastErr = err
			continue
		}
		for _, a := range addrs {
			if !usable(a) {
				continue
			}
			e := Entry{Addr: a.Unmap().String(), Source: src, Seen: now}
			got = append(got, e)
			if first == "" {
				first = src
			}
		}
	}
	if len(got) == 0 {
		d.Failures.Add(1)
		d.log().Warn("self_ip discovery failed", "err", lastErr)
		d.applyLocked()
		return fmt.Errorf("self_ip: no address discovered: %v", lastErr)
	}
	d.mu.Lock()
	d.entries = merge(d.entries, got, now, d.keep())
	d.lastOK = now
	d.lastSource = first
	d.mu.Unlock()
	d.applyLocked()
	return nil
}

// Start loads persist, runs one discovery with a bounded wait, applies, then ticks.
// The listener must not open until Start returns (the first RunOnce).
func (d *Discoverer) Start(parent context.Context) {
	if err := d.Load(); err != nil {
		d.log().Warn("self_ip load", "err", err)
	} else {
		d.applyLocked()
	}
	wait := d.Cfg.StartupWait
	if wait <= 0 {
		wait = DefaultStartupWait
	}
	ctx, cancel := context.WithTimeout(parent, wait)
	err := d.RunOnce(ctx)
	cancel()
	if err != nil {
		d.log().Warn("self_ip startup", "err", err)
	}
	if err := d.Persist(); err != nil {
		d.log().Warn("self_ip persist", "err", err)
	}
	go d.loop(parent)
}

func (d *Discoverer) loop(ctx context.Context) {
	t := time.NewTicker(d.interval())
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			c, cancel := context.WithTimeout(ctx, DefaultStartupWait)
			err := d.RunOnce(c)
			cancel()
			if err != nil {
				continue
			}
			if err := d.Persist(); err != nil {
				d.log().Warn("self_ip persist", "err", err)
			}
		}
	}
}

func (d *Discoverer) applyLocked() {
	if d.Guard == nil {
		return
	}
	d.mu.Lock()
	ents := append([]Entry(nil), d.entries...)
	d.mu.Unlock()
	var prefixes []netip.Prefix
	seen := map[string]struct{}{}
	for _, e := range ents {
		p, err := netip.ParsePrefix(prefixOf(e.Addr))
		if err != nil {
			continue
		}
		s := p.String()
		if _, ok := seen[s]; ok {
			continue
		}
		seen[s] = struct{}{}
		prefixes = append(prefixes, p)
	}
	d.Guard.SetSelfDeny(prefixes)
}

func (d *Discoverer) resolve(ctx context.Context, src string) ([]netip.Addr, error) {
	switch {
	case strings.HasPrefix(src, "ddns:"):
		host := strings.TrimPrefix(src, "ddns:")
		if host == "" {
			return nil, fmt.Errorf("empty ddns host")
		}
		lookup := d.Lookup
		if lookup == nil {
			lookup = func(ctx context.Context, h string) ([]netip.Addr, error) {
				return net.DefaultResolver.LookupNetIP(ctx, "ip", h)
			}
		}
		return lookup(ctx, host)
	case strings.HasPrefix(src, "trace:"):
		raw := strings.TrimPrefix(src, "trace:")
		if d.Get == nil {
			return nil, fmt.Errorf("no trace getter")
		}
		body, err := d.Get(ctx, raw)
		if err != nil {
			return nil, err
		}
		a, err := parseTrace(string(body))
		if err != nil {
			return nil, err
		}
		return []netip.Addr{a}, nil
	default:
		return nil, fmt.Errorf("unknown self_ip source %q", src)
	}
}

func parseTrace(body string) (netip.Addr, error) {
	for _, line := range strings.Split(body, "\n") {
		line = strings.TrimSpace(line)
		k, v, ok := strings.Cut(line, "=")
		if !ok || k != "ip" {
			continue
		}
		a, err := netip.ParseAddr(strings.TrimSpace(v))
		if err != nil {
			return netip.Addr{}, fmt.Errorf("trace ip=: %w", err)
		}
		return a.Unmap(), nil
	}
	return netip.Addr{}, fmt.Errorf("trace: no ip= line")
}

func usable(a netip.Addr) bool {
	if !a.IsValid() {
		return false
	}
	a = a.Unmap()
	if !a.IsGlobalUnicast() || a.IsPrivate() {
		return false
	}
	return true
}

func prefixOf(addr string) string {
	a, err := netip.ParseAddr(addr)
	if err != nil {
		return ""
	}
	a = a.Unmap()
	bits := 32
	if a.Is6() {
		bits = 64
	}
	p, err := a.Prefix(bits)
	if err != nil {
		return ""
	}
	return p.String()
}

func prune(ents []Entry, now time.Time, keep time.Duration) []Entry {
	out := ents[:0]
	for _, e := range ents {
		if keep > 0 && now.Sub(e.Seen) > keep {
			continue
		}
		out = append(out, e)
	}
	return append([]Entry(nil), out...)
}

func merge(old, got []Entry, now time.Time, keep time.Duration) []Entry {
	byAddr := map[string]Entry{}
	for _, e := range prune(old, now, keep) {
		byAddr[e.Addr] = e
	}
	for _, e := range got {
		byAddr[e.Addr] = e
	}
	out := make([]Entry, 0, len(byAddr))
	for _, e := range byAddr {
		out = append(out, e)
	}
	return out
}