package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"
	"time"

	"internetpressure.io/probe-agent/internal/batcher"
	checkdns "internetpressure.io/probe-agent/internal/checks/dns"
	checkhttp "internetpressure.io/probe-agent/internal/checks/http"
	checkping "internetpressure.io/probe-agent/internal/checks/ping"
	checktrace "internetpressure.io/probe-agent/internal/checks/traceroute"
	"internetpressure.io/probe-agent/internal/client"
	"internetpressure.io/probe-agent/internal/config"
	"internetpressure.io/probe-agent/internal/health"
	"internetpressure.io/probe-agent/internal/identity"
	"internetpressure.io/probe-agent/internal/protocol"
	"internetpressure.io/probe-agent/internal/sched"
	"internetpressure.io/probe-agent/internal/signer"
	"internetpressure.io/probe-agent/internal/spool"
	"internetpressure.io/probe-agent/internal/update"
)

const (
	cachedConfigFile   = "config.json"
	configBackoffMin   = 5 * time.Second
	configBackoffMax   = 5 * time.Minute
	defaultRefresh     = 300 * time.Second
	clockWarnThreshold = 120 * time.Second
	clockWarnEvery     = 10 * time.Minute
)

// agent wires every component together.
type agent struct {
	cfg *config.Config
	log *slog.Logger

	client   *client.Client
	spool    *spool.Spool
	health   *health.State
	batcher  *batcher.Batcher
	sched    *sched.Scheduler
	identity *identity.Service
	updater  *update.Updater

	httpCheck  *checkhttp.Checker
	dnsCheck   *checkdns.Checker
	pingCheck  *checkping.Checker
	traceCheck *checktrace.Checker
	caps       []string

	remote    atomic.Pointer[protocol.RemoteConfig]
	refreshCh chan struct{}
	lastClock time.Time
	clockMu   sync.Mutex
	updated   atomic.Bool
	cancel    context.CancelFunc
}

func newAgent(cfg *config.Config, log *slog.Logger) (*agent, error) {
	if err := os.MkdirAll(cfg.DataDir, 0o700); err != nil {
		return nil, fmt.Errorf("data_dir: %w", err)
	}
	sg, err := signer.New(cfg.ProbeID, cfg.Key)
	if err != nil {
		return nil, err
	}
	cl, err := client.New(cfg.IngestURL, sg, version)
	if err != nil {
		return nil, err
	}
	sp, err := spool.Open(filepath.Join(cfg.DataDir, "spool"), spool.DefaultCap)
	if err != nil {
		return nil, err
	}

	a := &agent{cfg: cfg, log: log, client: cl, spool: sp, refreshCh: make(chan struct{}, 1)}
	ua := client.UserAgent(version)
	a.httpCheck = &checkhttp.Checker{UserAgent: ua}
	a.dnsCheck = &checkdns.Checker{}
	a.pingCheck = &checkping.Checker{ICMP: checkping.Detect()}
	a.traceCheck = checktrace.Detect()
	a.caps = []string{"http", "dns"}
	if a.pingCheck.ICMP {
		a.caps = append(a.caps, "ping")
	} else {
		a.caps = append(a.caps, "tcp")
		log.Warn("unprivileged ICMP unavailable: ping checks fall back to tcp connects", "hint", checkping.LinuxHint())
	}
	if a.traceCheck.Available() {
		a.caps = append(a.caps, "traceroute")
	} else {
		log.Warn("traceroute binary not found: traceroute disabled")
	}

	a.identity = identity.New(ua)
	a.sched = sched.New(a.runJob, sched.Options{
		MaxConcurrency: cfg.MaxConcurrency,
		Logger:         log,
		CanPing:        a.pingCheck.ICMP,
		CanTraceroute:  a.traceCheck.Available(),
	})
	a.health = health.New(cfg.ProbeID, version, a.caps, health.Gauges{
		Buffered:      func() int { return a.batcher.Buffered() },
		SpoolBytes:    sp.Bytes,
		ClockOffsetMs: cl.Clock.OffsetMs,
		Targets:       a.sched.Targets,
		Identity:      a.identity.Current,
	})
	a.batcher = batcher.New(cfg.ProbeID, version, cl, sp, batcher.Hooks{
		Health:         a.health.Snapshot,
		OnResponse:     a.onBatchResponse,
		OnFlushFailure: a.health.RecordFlushFailure,
		OnDelivered: func(t time.Time) {
			a.health.SetLastFlush(t)
			a.checkClock()
		},
	}, log)
	a.updater = &update.Updater{Fetcher: cl, Version: version, DataDir: cfg.DataDir, Log: log}
	return a, nil
}

// Run starts every loop and blocks until ctx is cancelled (SIGINT/SIGTERM) or a self-update requires a restart.
func (a *agent) Run(parent context.Context) error {
	ctx, cancel := context.WithCancel(parent)
	a.cancel = cancel
	defer cancel()
	a.log.Info("ip-probe starting", "version", version, "probe_id", a.cfg.ProbeID, "ingest_url", a.cfg.IngestURL,
		"data_dir", a.cfg.DataDir, "listen", a.cfg.Listen, "capabilities", a.caps, "max_concurrency", a.cfg.MaxConcurrency)

	// Offline start: reuse the last good config so checks begin before the server answers.
	if rc, err := a.loadCachedConfig(); err == nil {
		a.log.Info("starting with cached config", "config_version", rc.ConfigVersion, "targets", len(rc.Targets))
		a.applyRemote(rc, false)
	}

	var wg sync.WaitGroup
	start := func(name string, f func(context.Context)) {
		wg.Add(1)
		go func() {
			defer wg.Done()
			f(ctx)
			a.log.Debug("loop stopped", "loop", name)
		}()
	}
	go func() {
		if err := a.health.Serve(ctx, a.cfg.Listen); err != nil {
			a.log.Error("local health endpoint failed", "listen", a.cfg.Listen, "err", err)
		}
	}()
	var schedDone sync.WaitGroup
	schedDone.Add(1)
	go func() { defer schedDone.Done(); a.sched.Run(ctx) }()
	start("batcher", a.batcher.Run)
	start("config", a.configLoop)
	start("identity", func(ctx context.Context) {
		a.identity.Run(ctx, func(id *protocol.Identity, err error) {
			if err != nil {
				a.log.Warn("identity lookup failed", "err", err)
				return
			}
			a.log.Info("identity", "public_ip", id.PublicIP, "asn", id.ASN, "org", id.Org, "country", id.Country,
				"city", id.City, "source", id.Source)
		})
	})
	if a.cfg.AllowSelfUpdate {
		start("update", func(ctx context.Context) {
			a.updater.Run(ctx, func() {
				a.updated.Store(true)
				a.log.Info("self-update installed; shutting down for the supervisor to restart")
				cancel()
			})
		})
	} else {
		a.log.Info("self-update disabled by configuration")
	}

	<-ctx.Done()
	a.log.Info("shutting down", "reason", reason(parent, a.updated.Load()))
	// 1. stop scheduling and wait for in-flight checks; 2. final flush (one 5 s attempt, then spool).
	schedDone.Wait()
	wg.Wait()
	a.batcher.Stop()
	a.log.Info("stopped", "buffered", a.batcher.Buffered(), "spool_bytes", a.spool.Bytes())
	return nil
}

func reason(parent context.Context, updated bool) string {
	switch {
	case updated:
		return "self-update"
	case parent.Err() != nil:
		return "signal"
	default:
		return "internal"
	}
}

// runJob executes one scheduled (target, family) job. Measurements from cancelled checks are dropped.
func (a *agent) runJob(ctx context.Context, job sched.Job) {
	switch job.Family {
	case sched.FamilyHTTP:
		a.emit(ctx, a.httpCheck.Run(ctx, job.Target))
	case sched.FamilyDNS:
		a.emit(ctx, a.dnsCheck.Run(ctx, job.Target, a.resolvers())...)
	case sched.FamilyPing:
		a.emit(ctx, a.pingCheck.Run(ctx, job.Target))
	case sched.FamilyTCP:
		a.emit(ctx, a.pingCheck.RunTCP(ctx, job.Target))
	case sched.FamilyTraceroute:
		tr, err := a.traceCheck.Run(ctx, job.Target)
		if ctx.Err() != nil {
			return
		}
		if err != nil {
			a.log.Debug("traceroute failed", "target", job.Target.TargetID, "err", err)
			return
		}
		a.batcher.AddTraceroute(tr)
	}
}

func (a *agent) emit(ctx context.Context, ms ...protocol.Measurement) {
	if ctx.Err() != nil {
		return
	}
	for _, m := range ms {
		a.health.Record(m)
	}
	a.batcher.Add(ms...)
}

// resolvers returns the DNS resolver list: local override, else server config, else defaults.
func (a *agent) resolvers() []protocol.Resolver {
	if len(a.cfg.ResolversOverride) > 0 {
		out := make([]protocol.Resolver, 0, len(a.cfg.ResolversOverride))
		for _, r := range a.cfg.ResolversOverride {
			id, addr, err := config.ParseResolver(r)
			if err == nil {
				out = append(out, protocol.Resolver{ID: id, Address: addr})
			}
		}
		return out
	}
	if rc := a.remote.Load(); rc != nil && len(rc.Resolvers) > 0 {
		return rc.Resolvers
	}
	return checkdns.DefaultResolvers
}

// configLoop fetches /config with backoff until it succeeds, then every config_refresh_seconds or on demand.
func (a *agent) configLoop(ctx context.Context) {
	backoff := configBackoffMin
	for {
		rc, err := a.client.GetConfig(ctx)
		if ctx.Err() != nil {
			return
		}
		var wait time.Duration
		if err != nil {
			a.log.Warn("config fetch failed", "err", err, "retry_in", backoff)
			wait = backoff
			backoff *= 2
			if backoff > configBackoffMax {
				backoff = configBackoffMax
			}
		} else {
			backoff = configBackoffMin
			a.applyRemote(rc, true)
			wait = defaultRefresh
			if rc.Schedule.ConfigRefreshSeconds > 0 {
				wait = time.Duration(rc.Schedule.ConfigRefreshSeconds) * time.Second
			}
			if wait < 30*time.Second {
				wait = 30 * time.Second
			}
		}
		select {
		case <-ctx.Done():
			return
		case <-time.After(wait):
		case <-a.refreshCh:
		}
	}
}

// applyRemote installs a remote config everywhere and (optionally) persists it for offline starts.
func (a *agent) applyRemote(rc *protocol.RemoteConfig, persist bool) {
	if rc.Probe.ProbeID != "" && rc.Probe.ProbeID != a.cfg.ProbeID {
		a.log.Warn("server config is for another probe id", "got", rc.Probe.ProbeID, "want", a.cfg.ProbeID)
	}
	prev := a.remote.Load()
	a.remote.Store(rc)
	a.sched.Apply(rc)
	a.sched.SetBoost(rc.Schedule.Boost)
	a.batcher.Configure(rc.Schedule.BatchFlushSeconds, rc.Schedule.MaxBatch)
	a.health.SetLastConfig(time.Now(), rc.ConfigVersion)
	if !rc.Probe.Enabled {
		a.log.Warn("probe disabled by server: idling (health only)")
	}
	if prev == nil || prev.ConfigVersion != rc.ConfigVersion {
		a.log.Info("config applied", "config_version", rc.ConfigVersion, "targets", len(rc.Targets),
			"resolvers", len(rc.Resolvers), "enabled", rc.Probe.Enabled, "tiers", rc.Schedule.Tiers,
			"dns_every", rc.Schedule.DNSEvery, "ping_every", rc.Schedule.PingEvery, "traceroute_every", rc.Schedule.TracerouteEvery)
	}
	if persist {
		if err := a.saveCachedConfig(rc); err != nil {
			a.log.Warn("could not persist config", "err", err)
		}
	}
	a.checkClock()
}

// onBatchResponse reacts to a batch reply: new config_version → refresh now; boost → apply.
func (a *agent) onBatchResponse(resp *protocol.BatchResponse) {
	if resp.Boost != nil {
		a.sched.SetBoost(resp.Boost)
	}
	cur := a.remote.Load()
	if resp.ConfigVersion != "" && (cur == nil || cur.ConfigVersion != resp.ConfigVersion) {
		a.log.Info("server announced a new config version", "config_version", resp.ConfigVersion)
		select {
		case a.refreshCh <- struct{}{}:
		default:
		}
	}
}

// checkClock warns (rate-limited) when the estimated clock offset exceeds 120 s.
func (a *agent) checkClock() {
	off := a.client.Clock.Offset()
	if off < 0 {
		off = -off
	}
	if off <= clockWarnThreshold {
		return
	}
	a.clockMu.Lock()
	defer a.clockMu.Unlock()
	if time.Since(a.lastClock) < clockWarnEvery {
		return
	}
	a.lastClock = time.Now()
	a.log.Warn("local clock differs from server by more than 120 s; fix NTP", "clock_offset_ms", a.client.Clock.OffsetMs())
}

func (a *agent) cachedConfigPath() string { return filepath.Join(a.cfg.DataDir, cachedConfigFile) }

func (a *agent) loadCachedConfig() (*protocol.RemoteConfig, error) {
	data, err := os.ReadFile(a.cachedConfigPath())
	if err != nil {
		return nil, err
	}
	var rc protocol.RemoteConfig
	if err := json.Unmarshal(data, &rc); err != nil {
		return nil, err
	}
	if rc.ConfigVersion == "" {
		return nil, errors.New("cached config has no version")
	}
	return &rc, nil
}

func (a *agent) saveCachedConfig(rc *protocol.RemoteConfig) error {
	data, err := json.Marshal(rc)
	if err != nil {
		return err
	}
	tmp := a.cachedConfigPath() + ".tmp"
	if err := os.WriteFile(tmp, data, 0o600); err != nil {
		return err
	}
	return os.Rename(tmp, a.cachedConfigPath())
}