// Package bamboo is the site agent. It only makes outbound connections: one
// HTTP long-poll to hammerhead for instructions, and one WireGuard tunnel
// (UDP) to each blacktip edge that forwards ports to this site. Inside the
// tunnel it accepts proxied connections and dials the real services.
package bamboo

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/netip"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"git.veillax.com/Veillax/bamboo/discover"
	"git.veillax.com/Veillax/reef/docker"
	"git.veillax.com/Veillax/reef/model"
	"git.veillax.com/Veillax/reef/tunnel"
	"git.veillax.com/Veillax/reef/wg"
)

type Config struct {
	Hammerhead string
	Token      string
	DataDir    string          // holds this agent's WireGuard key
	Docker     bool            // report local Docker containers so they can be picked as targets
	Discover   discover.Config // opt-in: report listening services found on this host and configured subnets as suggested targets

	// Version, AutoUpdate and UpdateNow work as in blacktip.Config.
	Version    string
	AutoUpdate bool
	UpdateNow  func(ctx context.Context, progress func(stage string, pct int, line string)) (newVersion string, err error)
}

type Bamboo struct {
	cfg   Config
	http  *http.Client
	ctx   context.Context
	wgKey string
	wgPub string

	mu      sync.Mutex
	version int64
	tunIP   string
	tun     *wg.Device
	edges   []model.EdgeRef
	lastRep string

	updating  atomic.Bool
	updateMsg string // outcome of the update Hammerhead asked for, reported back
	updateID  string // the request that updateMsg answers
	updStage  string // progress of the running update, reported to Hammerhead
	updPct    int
	updLog    []string
	updSent   time.Time

	disc *discover.Discoverer // nil unless discovery is switched on

	mail    map[string]*mailRun
	mailErr map[string]string // mail egress ID -> "" (listening) or an error
}

func New(cfg Config) (*Bamboo, error) {
	if cfg.DataDir == "" {
		cfg.DataDir = "./bamboo-data"
	}
	key, err := loadOrCreateKey(filepath.Join(cfg.DataDir, "wg.key"))
	if err != nil {
		return nil, err
	}
	pub, err := wg.PublicKey(key)
	if err != nil {
		return nil, err
	}
	b := &Bamboo{cfg: cfg, http: &http.Client{Timeout: 40 * time.Second}, ctx: context.Background(), wgKey: key, wgPub: pub,
		mail: map[string]*mailRun{}, mailErr: map[string]string{}}
	if cfg.Discover.Enabled() {
		if b.disc, err = discover.New(cfg.Discover); err != nil {
			return nil, err
		}
	}
	return b, nil
}

// loadOrCreateKey keeps this agent's WireGuard identity across restarts.
func loadOrCreateKey(path string) (string, error) {
	if b, err := os.ReadFile(path); err == nil {
		if k := strings.TrimSpace(string(b)); k != "" {
			return k, nil
		}
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return "", err
	}
	k, err := wg.GenerateKey()
	if err != nil {
		return "", err
	}
	return k, os.WriteFile(path, []byte(k+"\n"), 0o600)
}

func (b *Bamboo) Run(ctx context.Context) error {
	b.ctx = ctx
	log.Printf("bamboo started; WireGuard public key %s", b.wgPub)
	go b.reportLoop(ctx)
	go b.refreshLoop(ctx)
	if b.disc != nil {
		go b.disc.Run(ctx, log.Printf)
	}
	for ctx.Err() == nil {
		cfg, err := b.fetchConfig(ctx)
		if err != nil {
			log.Printf("config: %v", err)
			select {
			case <-time.After(3 * time.Second):
			case <-ctx.Done():
			}
			continue
		}
		if cfg != nil {
			b.apply(cfg)
		}
	}
	b.mu.Lock()
	for _, r := range b.mail {
		r.ln.Close()
	}
	if b.tun != nil {
		b.tun.Close()
	}
	b.mu.Unlock()
	return nil
}

func (b *Bamboo) fetchConfig(ctx context.Context) (*model.SiteConfig, error) {
	b.mu.Lock()
	v := b.version
	b.mu.Unlock()
	req, _ := http.NewRequestWithContext(ctx, "GET", fmt.Sprintf("%s/site/config?v=%d", b.cfg.Hammerhead, v), nil)
	req.Header.Set("Authorization", "Bearer "+b.cfg.Token)
	resp, err := b.http.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode == http.StatusNoContent {
		return nil, nil
	}
	if resp.StatusCode != 200 {
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
		return nil, fmt.Errorf("hammerhead said %d: %s", resp.StatusCode, bytes.TrimSpace(body))
	}
	var cfg model.SiteConfig
	return &cfg, json.NewDecoder(resp.Body).Decode(&cfg)
}

func (b *Bamboo) apply(cfg *model.SiteConfig) {
	b.mu.Lock()
	b.version = cfg.Version
	b.edges = cfg.Edges
	if cfg.UpdateID != "" {
		b.requestUpdate(cfg.UpdateID)
	}
	if cfg.TunnelIP != b.tunIP && b.tun != nil {
		b.tun.Close()
		b.tun = nil
	}
	if b.tun == nil && cfg.TunnelIP != "" {
		if err := b.startTunnel(cfg.TunnelIP); err != nil {
			log.Printf("start WireGuard: %v", err)
		}
	}
	b.mu.Unlock()
	b.syncPeers()
	b.syncMail(cfg.MailEgress)
	go b.report()
}

// startTunnel creates the WireGuard device and starts serving inside it. Caller holds b.mu.
func (b *Bamboo) startTunnel(ip string) error {
	addr, err := netip.ParseAddr(ip)
	if err != nil {
		return err
	}
	dev, err := wg.New(b.wgKey, addr, 0) // we dial out, so any local UDP port will do
	if err != nil {
		return err
	}
	tcp, err := dev.ListenTCP(tunnel.ProxyTCPPort)
	if err != nil {
		dev.Close()
		return err
	}
	udp, err := dev.ListenUDP(tunnel.ProxyUDPPort)
	if err != nil {
		dev.Close()
		return err
	}
	b.tun, b.tunIP = dev, ip
	go serveTCP(tcp)
	go serveUDP(udp)
	log.Printf("WireGuard up at %s", ip)
	return nil
}

// syncPeers makes the tunnel match the edges hammerhead lists. Unchanged peers keep their sessions.
func (b *Bamboo) syncPeers() {
	b.mu.Lock()
	tun, edges := b.tun, append([]model.EdgeRef(nil), b.edges...)
	b.mu.Unlock()
	if tun == nil {
		return
	}
	var peers []wg.Peer
	for _, e := range edges {
		if e.PublicKey == "" || e.IP == "" {
			continue // that edge hasn't started yet
		}
		peers = append(peers, wg.Peer{PublicKey: e.PublicKey, AllowedIP: e.IP, Endpoint: e.Endpoint})
	}
	if err := tun.SetPeers(peers); err != nil {
		log.Printf("peers: %v", err)
	}
}

// refreshLoop re-resolves edge hostnames, so a changed DNS record is followed.
func (b *Bamboo) refreshLoop(ctx context.Context) {
	t := time.NewTicker(60 * time.Second)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			b.syncPeers()
		}
	}
}

// connected lists the edges with a recent WireGuard handshake.
func (b *Bamboo) connected() []string {
	b.mu.Lock()
	tun, edges := b.tun, b.edges
	b.mu.Unlock()
	ids := []string{}
	if tun == nil {
		return ids
	}
	hs := tun.Handshakes()
	for _, e := range edges {
		if t, ok := hs[e.PublicKey]; ok && time.Since(t) < 3*time.Minute {
			ids = append(ids, e.ID)
		}
	}
	sort.Strings(ids)
	return ids
}

func (b *Bamboo) report() {
	ids := b.connected()
	var containers []model.Container
	if b.cfg.Docker {
		ctx, cancel := context.WithTimeout(b.ctx, 5*time.Second)
		containers, _ = docker.List(ctx)
		cancel()
	}
	var found []model.DiscoveredService
	if b.disc != nil {
		found = b.disc.Results()
	}
	b.mu.Lock()
	msg, uid, stage, pct, ulog := b.updateMsg, b.updateID, b.updStage, b.updPct, append([]string(nil), b.updLog...)
	b.mu.Unlock()
	body, _ := json.Marshal(model.SiteReport{PublicKey: b.wgPub, Connected: ids, Containers: containers, Discovered: found, MailEgress: b.mailStatus(),
		Version: b.cfg.Version, AutoUpdate: b.cfg.AutoUpdate, UpdateMsg: msg, UpdateID: uid, UpdateStage: stage, UpdatePct: pct, UpdateLog: ulog})
	req, _ := http.NewRequest("POST", b.cfg.Hammerhead+"/site/report", bytes.NewReader(body))
	req.Header.Set("Authorization", "Bearer "+b.cfg.Token)
	req.Header.Set("Content-Type", "application/json")
	resp, err := b.http.Do(req)
	if err != nil {
		return
	}
	resp.Body.Close()
}

// reportLoop reports every 20 seconds, and straight away when the set of connected edges changes.
func (b *Bamboo) reportLoop(ctx context.Context) {
	b.report()
	slow := time.NewTicker(20 * time.Second)
	fast := time.NewTicker(3 * time.Second)
	defer slow.Stop()
	defer fast.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-slow.C:
			b.report()
		case <-fast.C:
			sig := strings.Join(b.connected(), ",")
			b.mu.Lock()
			changed := sig != b.lastRep
			b.lastRep = sig
			b.mu.Unlock()
			if changed {
				b.report()
			}
		}
	}
}

// ---------- serving proxied connections ----------

func serveTCP(ln net.Listener) {
	for {
		c, err := ln.Accept()
		if err != nil {
			return
		}
		go handleStream(c)
	}
}

func handleStream(st net.Conn) {
	st.SetReadDeadline(time.Now().Add(10 * time.Second))
	br := bufio.NewReader(st)
	h, err := tunnel.ReadHeader(br)
	st.SetReadDeadline(time.Time{})
	if err != nil || h.Proto != "tcp" {
		st.Close()
		return
	}
	c, err := net.DialTimeout("tcp", net.JoinHostPort(h.Host, strconv.Itoa(h.Port)), 8*time.Second)
	if err != nil {
		st.Close() // the edge reads this as "target unreachable"
		return
	}
	tunnel.Proxy(st, br, c, c)
}

type udpSession struct {
	conn net.Conn
	last time.Time
}

// serveUDP relays datagrams between the tunnel and UDP targets. Each datagram
// starts with a type byte: 0 = target header + payload (creates the session),
// 1 = payload for an existing session. Sessions are keyed by the edge's source port.
func serveUDP(pc net.PacketConn) {
	var mu sync.Mutex
	sessions := map[string]*udpSession{}
	go func() {
		t := time.NewTicker(15 * time.Second)
		defer t.Stop()
		for range t.C {
			mu.Lock()
			for _, s := range sessions {
				if time.Since(s.last) > 60*time.Second {
					s.conn.Close()
				}
			}
			mu.Unlock()
		}
	}()
	buf := make([]byte, 65535)
	for {
		n, from, err := pc.ReadFrom(buf)
		if err != nil {
			return
		}
		if n < 2 {
			continue
		}
		key := from.String()
		data := buf[:n]
		var payload []byte
		mu.Lock()
		s := sessions[key]
		switch data[0] {
		case 0:
			if len(data) < 3 {
				mu.Unlock()
				continue
			}
			hl := int(data[1])<<8 | int(data[2])
			if len(data) < 3+hl {
				mu.Unlock()
				continue
			}
			payload = data[3+hl:]
			if s == nil {
				var h tunnel.Header
				if json.Unmarshal(data[3:3+hl], &h) != nil || h.Proto != "udp" {
					mu.Unlock()
					continue
				}
				c, err := net.Dial("udp", net.JoinHostPort(h.Host, strconv.Itoa(h.Port)))
				if err != nil {
					mu.Unlock()
					continue
				}
				s = &udpSession{conn: c}
				sessions[key] = s
				go func(s *udpSession, from net.Addr, key string) {
					defer func() {
						s.conn.Close()
						mu.Lock()
						if sessions[key] == s {
							delete(sessions, key)
						}
						mu.Unlock()
					}()
					rb := make([]byte, 65535)
					for {
						n, err := s.conn.Read(rb)
						if err != nil {
							return
						}
						mu.Lock()
						s.last = time.Now()
						mu.Unlock()
						if _, err := pc.WriteTo(rb[:n], from); err != nil {
							return
						}
					}
				}(s, from, key)
			}
		case 1:
			payload = data[1:]
		default:
			mu.Unlock()
			continue
		}
		if s != nil {
			s.last = time.Now()
		}
		mu.Unlock()
		if s != nil && len(payload) > 0 {
			s.conn.Write(payload)
		}
	}
}

// requestUpdate acts on "update now" from Hammerhead (id identifies the request). It only ever installs
// what this machine's own auto-update would (the newest release from its configured server), and only
// when auto-update is on. Caller holds b.mu.
func (b *Bamboo) requestUpdate(id string) {
	if id == b.updateID {
		return // already handled, or in progress
	}
	b.updateID, b.updateMsg = id, ""
	b.updStage, b.updPct, b.updLog = "starting", 0, nil
	b.updNote("Update requested from the dashboard")
	if !b.cfg.AutoUpdate || b.cfg.UpdateNow == nil {
		b.updateMsg = "auto-update is off on this machine (set BAMBOO_AUTO_UPDATE=1 and restart it)"
		b.updStage = "failed"
		b.updNote("Refused: " + b.updateMsg)
		go b.report()
		return
	}
	if !b.updating.CompareAndSwap(false, true) {
		b.updateMsg = "another update is already running"
		b.updStage = "failed"
		b.updNote("Refused: " + b.updateMsg)
		go b.report()
		return
	}
	go func() {
		defer b.updating.Store(false)
		tag, err := b.cfg.UpdateNow(b.ctx, b.updProgress) // on success this restarts the process and does not return
		b.mu.Lock()
		switch {
		case err != nil:
			b.updateMsg = "update failed: " + err.Error()
		case tag == "":
			b.updateMsg = "already on the newest release"
		default:
			b.updateMsg = "updated to " + tag + ", restarting"
		}
		b.mu.Unlock()
		b.report()
	}()
}

// updNote adds a timestamped line to the update log. Caller holds b.mu.
func (b *Bamboo) updNote(line string) {
	b.updLog = append(b.updLog, time.Now().Format("15:04:05")+" "+line)
	if len(b.updLog) > 200 {
		b.updLog = b.updLog[len(b.updLog)-200:]
	}
}

// updProgress is called by the updater as it works; Hammerhead gets a fresh report at most once a
// second (and straight away for every new log line), so its progress bar follows along.
func (b *Bamboo) updProgress(stage string, pct int, line string) {
	b.mu.Lock()
	b.updStage = stage
	if pct >= 0 {
		b.updPct = pct
	}
	if line != "" {
		b.updNote(line)
	}
	send := line != "" || time.Since(b.updSent) > time.Second
	if send {
		b.updSent = time.Now()
	}
	b.mu.Unlock()
	if send {
		go b.report()
	}
}