// Package traceroute runs the system traceroute binary (default UDP probes, which works unprivileged on macOS)
// and parses its output into the protocol's Traceroute record.
//
//	traceroute -n -q 1 -w 2 -m 30 <ip>      (60 s process timeout)
//
// route_hash = sha1 hex over the "|"-joined hop IPs, unanswered hops kept as "*".
package traceroute

import (
	"bytes"
	"context"
	"crypto/sha1"
	"encoding/hex"
	"errors"
	"net"
	"os"
	"os/exec"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"time"

	"internetpressure.io/probe-agent/internal/protocol"
)

// ProcessTimeout bounds one traceroute run.
const ProcessTimeout = 60 * time.Second

// MaxHops passed as -m.
const MaxHops = 30

// Checker runs traceroutes with the detected binary. Binary is "" when unavailable.
type Checker struct {
	Binary  string
	Binary6 string // traceroute6 on macOS (optional)
}

// Detect locates the traceroute binary for this OS. Returns a Checker with Binary "" when none is found.
func Detect() *Checker {
	c := &Checker{}
	var candidates []string
	switch runtime.GOOS {
	case "darwin":
		candidates = []string{"/usr/sbin/traceroute"}
	default:
		candidates = []string{"/usr/bin/traceroute", "/usr/sbin/traceroute", "/bin/traceroute"}
	}
	for _, p := range candidates {
		if fi, err := os.Stat(p); err == nil && !fi.IsDir() {
			c.Binary = p
			break
		}
	}
	if c.Binary == "" {
		if p, err := exec.LookPath("traceroute"); err == nil {
			c.Binary = p
		}
	}
	if runtime.GOOS == "darwin" {
		if fi, err := os.Stat("/usr/sbin/traceroute6"); err == nil && !fi.IsDir() {
			c.Binary6 = "/usr/sbin/traceroute6"
		}
	}
	return c
}

// Available reports whether traceroute can run.
func (c *Checker) Available() bool { return c != nil && c.Binary != "" }

// Run resolves the target (pinned IP or first IPv4 of the hostname) and runs one traceroute.
func (c *Checker) Run(ctx context.Context, t protocol.Target) (protocol.Traceroute, error) {
	tr := protocol.Traceroute{TS: protocol.FormatTime(time.Now()), TargetID: t.TargetID, Hops: []protocol.Hop{}}
	if !c.Available() {
		return tr, errors.New("traceroute binary not available")
	}
	ip, err := resolve(ctx, t)
	if err != nil {
		return tr, err
	}
	tr.DestIP = ip.String()

	ctx, cancel := context.WithTimeout(ctx, ProcessTimeout)
	defer cancel()
	bin, args := c.command(ip)
	cmd := exec.CommandContext(ctx, bin, args...)
	cmd.WaitDelay = 2 * time.Second
	var out, stderr bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &stderr
	runErr := cmd.Run()
	// Partial output is still useful when the process was killed at the deadline.
	if out.Len() == 0 && runErr != nil {
		msg := strings.TrimSpace(stderr.String())
		if msg == "" {
			msg = runErr.Error()
		}
		return tr, errors.New("traceroute: " + msg)
	}
	hops := Parse(out.String())
	Fill(&tr, hops)
	return tr, nil
}

func (c *Checker) command(ip net.IP) (string, []string) {
	args := []string{"-n", "-q", "1", "-w", "2", "-m", strconv.Itoa(MaxHops)}
	if ip.To4() == nil {
		if c.Binary6 != "" {
			return c.Binary6, append(args, ip.String())
		}
		return c.Binary, append([]string{"-6"}, append(args, ip.String())...)
	}
	return c.Binary, append(args, ip.String())
}

func resolve(ctx context.Context, t protocol.Target) (net.IP, error) {
	if s := t.FixedIP(); s != "" {
		if ip := net.ParseIP(s); ip != nil {
			return ip, nil
		}
		return nil, errors.New("invalid pinned ip")
	}
	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
	defer cancel()
	addrs, err := net.DefaultResolver.LookupIPAddr(ctx, t.Hostname)
	if err != nil {
		return nil, err
	}
	for _, a := range addrs {
		if a.IP.To4() != nil {
			return a.IP, nil
		}
	}
	if len(addrs) > 0 {
		return addrs[0].IP, nil
	}
	return nil, errors.New("no address")
}

// hopLine matches " 3  10.0.0.1  12.345 ms" and " 2  *" (with -q 1 there is one probe per line). Annotations
// such as " !H" or " !X" after the RTT are ignored. On Linux an unanswered hop may print as " 2  * " too.
var hopLine = regexp.MustCompile(`^\s*(\d+)\s+(\*|[0-9a-fA-F.:]+)(?:\s+(\d+(?:\.\d+)?)\s*ms)?`)

// Parse converts raw traceroute stdout into hops, ordered by hop number. Lines that are not hops (the header
// line, warnings) are skipped. Missing hop numbers are filled as unanswered so the route hash is positional.
func Parse(out string) []protocol.Hop {
	byN := map[int]protocol.Hop{}
	maxN := 0
	for _, line := range strings.Split(out, "\n") {
		m := hopLine.FindStringSubmatch(line)
		if m == nil {
			continue
		}
		n, _ := strconv.Atoi(m[1])
		if n <= 0 || n > MaxHops {
			continue
		}
		h := protocol.Hop{N: n, IP: m[2]}
		if m[2] == "*" {
			h.RTTMs = nil
		} else if m[3] != "" {
			if v, err := strconv.ParseFloat(m[3], 64); err == nil {
				h.RTTMs = protocol.F(v)
			}
		}
		// If the same hop number appears twice (unusual with -q 1), keep the answered one.
		if prev, ok := byN[n]; ok && prev.IP != "*" && h.IP == "*" {
			continue
		}
		byN[n] = h
		if n > maxN {
			maxN = n
		}
	}
	hops := make([]protocol.Hop, 0, maxN)
	for n := 1; n <= maxN; n++ {
		if h, ok := byN[n]; ok {
			hops = append(hops, h)
		} else {
			hops = append(hops, protocol.Hop{N: n, IP: "*"})
		}
	}
	return hops
}

// Fill computes hop_count, reached, total_ms and route_hash from hops.
func Fill(tr *protocol.Traceroute, hops []protocol.Hop) {
	tr.Hops = hops
	tr.HopCount = len(hops)
	ips := make([]string, len(hops))
	for i, h := range hops {
		ips[i] = h.IP
	}
	tr.RouteHash = RouteHash(ips)
	tr.Reached = len(hops) > 0 && hops[len(hops)-1].IP == tr.DestIP
	tr.TotalMs = nil
	for i := len(hops) - 1; i >= 0; i-- {
		if hops[i].RTTMs != nil {
			tr.TotalMs = protocol.F(*hops[i].RTTMs)
			break
		}
	}
}

// RouteHash is sha1 hex over "ip1|ip2|*|ip4…".
func RouteHash(ips []string) string {
	sum := sha1.Sum([]byte(strings.Join(ips, "|")))
	return hex.EncodeToString(sum[:])
}