// Package diagnostics runs bounded, read-only target diagnostics without a
// shell. It never accepts a command name or arbitrary argument list from API
// clients.
package diagnostics

import (
	"context"
	"errors"
	"fmt"
	"net"
	"os/exec"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	KindNetworkTest               = "network_test"
	KindTraceroute                = "traceroute"
	KindTCPConnectionTest         = "tcp_connection_test"
	DefaultTracerouteMaxHops      = 10
	MaximumTracerouteMaxHops      = 64
	TracerouteProbeTimeoutSeconds = 3
	maxTimeout                    = 300
	maxOutput                     = 1024 * 1024
)

type Request struct {
	Kind           string
	Target         string
	Port           int
	Count          int
	MaxHops        int
	TimeoutSeconds int
	MaxOutputBytes int
}

type Result struct {
	Kind            string
	Target          string
	Port            int
	Count           int
	MaxHops         int
	ExitCode        int
	Output          string
	OutputBytes     int
	OutputTruncated bool
	Reachable       *bool
}

func (r Result) OperationResult(timeoutSeconds, maxOutputBytes int) map[string]any {
	result := map[string]any{
		"kind":             r.Kind,
		"target":           r.Target,
		"timeout_seconds":  timeoutSeconds,
		"max_output_bytes": maxOutputBytes,
		"exit_code":        r.ExitCode,
		"output":           r.Output,
		"output_bytes":     r.OutputBytes,
		"output_truncated": r.OutputTruncated,
		"adapter":          "process",
	}
	if r.Reachable != nil {
		result["reachable"] = *r.Reachable
	}
	if r.Port != 0 {
		result["port"] = r.Port
		result["adapter"] = "socket"
	}
	if r.Count != 0 {
		result["count"] = r.Count
	}
	if r.MaxHops != 0 {
		result["max_hops"] = r.MaxHops
	}
	return result
}

func TracerouteOperationTimeoutSeconds(maxHops int) int {
	if maxHops < 1 || maxHops > MaximumTracerouteMaxHops {
		maxHops = DefaultTracerouteMaxHops
	}
	return maxHops*TracerouteProbeTimeoutSeconds + 5
}

type Runner interface {
	Run(context.Context, Request) (Result, error)
}

type RunnerFunc func(context.Context, Request) (Result, error)

func (f RunnerFunc) Run(ctx context.Context, request Request) (Result, error) {
	return f(ctx, request)
}

type Execution struct {
	Output    string
	Bytes     int
	Truncated bool
	ExitCode  int
}

type OutputProgress struct {
	Output    string
	Bytes     int
	Truncated bool
}

type outputReporterContextKey struct{}

func WithOutputReporter(ctx context.Context, reporter func(OutputProgress)) context.Context {
	if reporter == nil {
		return ctx
	}
	return context.WithValue(ctx, outputReporterContextKey{}, reporter)
}

func ReportOutput(ctx context.Context, progress OutputProgress) {
	reporter, _ := ctx.Value(outputReporterContextKey{}).(func(OutputProgress))
	if reporter != nil {
		reporter(progress)
	}
}

type Executor interface {
	Execute(context.Context, string, []string, int) (Execution, error)
}

type ProcessRunnerOptions struct {
	PingPath       string
	TraceroutePath string
	Executor       Executor
}

type ProcessRunner struct {
	pingPath       string
	traceroutePath string
	executor       Executor
}

func NewProcessRunner(options ProcessRunnerOptions) (*ProcessRunner, error) {
	if options.PingPath == "" {
		options.PingPath = "/bin/ping"
	}
	if options.TraceroutePath == "" {
		options.TraceroutePath = "/usr/bin/traceroute"
	}
	if !isAbsoluteExecutablePath(options.PingPath) || !isAbsoluteExecutablePath(options.TraceroutePath) {
		return nil, errors.New("diagnostic executable paths must be absolute")
	}
	if options.Executor == nil {
		options.Executor = OSExecutor{}
	}
	return &ProcessRunner{
		pingPath:       cleanExecutablePath(options.PingPath),
		traceroutePath: cleanExecutablePath(options.TraceroutePath),
		executor:       options.Executor,
	}, nil
}

func (r *ProcessRunner) Run(ctx context.Context, request Request) (Result, error) {
	if err := ValidateRequest(request); err != nil {
		return Result{}, err
	}
	if request.Kind == KindTCPConnectionTest {
		return runTCPConnectionTest(ctx, request)
	}
	if request.Kind == KindTraceroute {
		return r.runPingTraceroute(ctx, request)
	}
	pingCount := request.Count
	if pingCount == 0 {
		pingCount = 4
	}
	path := r.pingPath
	args := []string{"-c", strconv.Itoa(pingCount), "-s", "56", "-W", "1", "-w", strconv.Itoa(pingCount + 1), request.Target}
	executionContext, cancel := context.WithTimeout(ctx, time.Duration(request.TimeoutSeconds)*time.Second)
	defer cancel()
	execution, err := r.executor.Execute(executionContext, path, args, request.MaxOutputBytes)
	if err != nil {
		return Result{}, err
	}
	result := Result{
		Kind:            request.Kind,
		Target:          request.Target,
		Count:           mapPingCount(request.Kind, pingCount),
		ExitCode:        execution.ExitCode,
		Output:          execution.Output,
		OutputBytes:     execution.Bytes,
		OutputTruncated: execution.Truncated,
	}
	if request.Kind == KindNetworkTest {
		reachable := execution.ExitCode == 0
		result.Reachable = &reachable
	}
	return result, nil
}

func (r *ProcessRunner) runPingTraceroute(ctx context.Context, request Request) (Result, error) {
	maxHops := request.MaxHops
	if maxHops == 0 {
		maxHops = DefaultTracerouteMaxHops
	}
	executionContext, cancel := context.WithTimeout(ctx, time.Duration(request.TimeoutSeconds)*time.Second)
	defer cancel()

	lines := []string{fmt.Sprintf("traceroute to %s, %d hops max, ICMP ping probes", request.Target, maxHops)}
	reachable := false
	exitCode := 1
	truncated := false
	for ttl := 1; ttl <= maxHops; ttl++ {
		probeContext := context.WithValue(executionContext, outputReporterContextKey{}, (func(OutputProgress))(nil))
		execution, err := r.executor.Execute(probeContext, r.pingPath, []string{
			"-c", "1", "-s", "56", "-W", "1", "-t", strconv.Itoa(ttl), request.Target,
		}, request.MaxOutputBytes)
		if err != nil {
			return Result{}, err
		}

		hop, reached, unreachable := pingProbeHop(request.Target, execution)
		lines = append(lines, fmt.Sprintf("%2d  %s", ttl, hop))
		output, outputTruncated := boundedOutput(strings.Join(lines, "\n"), request.MaxOutputBytes)
		truncated = truncated || outputTruncated || execution.Truncated
		ReportOutput(executionContext, OutputProgress{Output: output, Bytes: len(output), Truncated: truncated})
		if reached {
			reachable = true
			exitCode = 0
			break
		}
		if unreachable {
			break
		}
	}

	output, outputTruncated := boundedOutput(strings.Join(lines, "\n"), request.MaxOutputBytes)
	truncated = truncated || outputTruncated
	return Result{
		Kind:            request.Kind,
		Target:          request.Target,
		MaxHops:         maxHops,
		ExitCode:        exitCode,
		Output:          output,
		OutputBytes:     len(output),
		OutputTruncated: truncated,
		Reachable:       &reachable,
	}, nil
}

func pingProbeHop(target string, execution Execution) (hop string, reached bool, unreachable bool) {
	for _, line := range strings.Split(strings.ToValidUTF8(execution.Output, ""), "\n") {
		trimmed := strings.TrimSpace(line)
		lower := strings.ToLower(trimmed)
		if index := strings.Index(lower, " bytes from "); index >= 0 {
			detail := trimmed[index+len(" bytes from "):]
			if marker := strings.Index(strings.ToLower(detail), " icmp_seq="); marker >= 0 {
				detail = strings.TrimSuffix(strings.TrimSpace(detail[:marker]), ":")
			}
			if latency := pingLatency(trimmed); latency != "" {
				detail += "  " + latency
			}
			return detail, true, false
		}
		if strings.HasPrefix(lower, "from ") {
			detail := strings.TrimSpace(trimmed[len("From "):])
			if marker := strings.Index(strings.ToLower(detail), " icmp_seq="); marker >= 0 {
				detail = strings.TrimSpace(detail[:marker])
			}
			if strings.Contains(lower, "unreachable") || strings.Contains(lower, "prohibited") {
				return detail + "  !H", false, true
			}
			if strings.Contains(lower, "time to live exceeded") || strings.Contains(lower, "ttl expired") {
				return detail, false, false
			}
		}
	}
	if execution.ExitCode == 0 {
		return target, true, false
	}
	return "*", false, false
}

func pingLatency(line string) string {
	lower := strings.ToLower(line)
	for _, marker := range []string{"time=", "time<"} {
		if index := strings.Index(lower, marker); index >= 0 {
			value := strings.Fields(line[index+len(marker):])
			if len(value) > 0 {
				prefix := ""
				if marker == "time<" {
					prefix = "<"
				}
				return prefix + strings.TrimSuffix(value[0], "ms") + " ms"
			}
		}
	}
	return ""
}

func boundedOutput(output string, limit int) (string, bool) {
	if len(output) <= limit {
		return output, false
	}
	return strings.ToValidUTF8(output[:limit], ""), true
}

func tracerouteReached(target, output string, exitCode int) bool {
	if exitCode != 0 {
		return false
	}
	lines := strings.Split(strings.TrimSpace(output), "\n")
	for index := len(lines) - 1; index >= 0; index-- {
		line := strings.TrimSpace(lines[index])
		if line == "" {
			continue
		}
		lower := strings.ToLower(line)
		if strings.Contains(lower, "unreachable") || strings.Contains(lower, "no route") || strings.Contains(line, " !") {
			return false
		}
		if net.ParseIP(target) != nil {
			for _, field := range strings.Fields(line) {
				if strings.Trim(field, "()") == target {
					return true
				}
			}
			return false
		}
		return !strings.Contains(line, "*")
	}
	return false
}

func runTCPConnectionTest(ctx context.Context, request Request) (Result, error) {
	executionContext, cancel := context.WithTimeout(ctx, time.Duration(request.TimeoutSeconds)*time.Second)
	defer cancel()
	connection, err := (&net.Dialer{}).DialContext(executionContext, "tcp", net.JoinHostPort(request.Target, strconv.Itoa(request.Port)))
	connected := err == nil
	output := "tcp connection established"
	exitCode := 0
	if err != nil {
		output = strings.ToValidUTF8(err.Error(), "")
		exitCode = 1
	} else {
		_ = connection.Close()
	}
	outputBytes := len(output)
	truncated := outputBytes > request.MaxOutputBytes
	if len(output) > request.MaxOutputBytes {
		output = output[:request.MaxOutputBytes]
	}
	return Result{
		Kind: request.Kind, Target: request.Target, Port: request.Port, ExitCode: exitCode,
		Output: output, OutputBytes: len(output), OutputTruncated: truncated,
		Reachable: &connected,
	}, nil
}

func ValidateRequest(request Request) error {
	if request.Kind != KindNetworkTest && request.Kind != KindTraceroute && request.Kind != KindTCPConnectionTest {
		return errors.New("diagnostic kind must be network_test, traceroute, or tcp_connection_test")
	}
	if err := ValidateTarget(request.Target); err != nil {
		return err
	}
	if request.TimeoutSeconds < 1 || request.TimeoutSeconds > maxTimeout {
		return fmt.Errorf("diagnostic timeout_seconds must be between 1 and %d", maxTimeout)
	}
	if request.MaxOutputBytes < 1024 || request.MaxOutputBytes > maxOutput {
		return fmt.Errorf("diagnostic max_output_bytes must be between 1024 and %d", maxOutput)
	}
	if request.Kind == KindTCPConnectionTest && (request.Port < 1 || request.Port > 65535) {
		return errors.New("diagnostic port must be between 1 and 65535 for tcp_connection_test")
	}
	if request.Kind == KindNetworkTest && request.Count != 0 && (request.Count < 1 || request.Count > 20) {
		return errors.New("diagnostic count must be between 1 and 20 for network_test")
	}
	if request.Kind == KindTraceroute && request.MaxHops != 0 && (request.MaxHops < 1 || request.MaxHops > MaximumTracerouteMaxHops) {
		return fmt.Errorf("diagnostic max_hops must be between 1 and %d for traceroute", MaximumTracerouteMaxHops)
	}
	if request.Kind != KindTraceroute && request.MaxHops != 0 {
		return errors.New("diagnostic max_hops is only valid for traceroute")
	}
	return nil
}

func mapPingCount(kind string, count int) int {
	if kind != KindNetworkTest {
		return 0
	}
	return count
}

func ValidateTarget(target string) error {
	if target == "" {
		return errors.New("diagnostic target must not be empty")
	}
	if len(target) > 128 {
		return errors.New("diagnostic target must be 128 characters or fewer")
	}
	if strings.HasPrefix(target, "-") {
		return errors.New("diagnostic target must not begin with a command option")
	}
	for _, char := range target {
		if (char >= 'A' && char <= 'Z') || (char >= 'a' && char <= 'z') || (char >= '0' && char <= '9') || char == '.' || char == '-' || char == '_' || char == ':' {
			continue
		}
		return errors.New("diagnostic target contains invalid characters")
	}
	return nil
}

type OSExecutor struct{}

func (OSExecutor) Execute(ctx context.Context, path string, args []string, maxOutputBytes int) (Execution, error) {
	if !isAbsoluteExecutablePath(path) {
		return Execution{}, errors.New("diagnostic executable path must be absolute")
	}
	if maxOutputBytes <= 0 || maxOutputBytes > maxOutput {
		return Execution{}, errors.New("diagnostic output limit is invalid")
	}
	command := exec.CommandContext(ctx, path, args...)
	buffer := &limitedBuffer{
		limit: maxOutputBytes,
		reporter: func(progress OutputProgress) {
			ReportOutput(ctx, progress)
		},
	}
	command.Stdout = buffer
	command.Stderr = buffer
	err := command.Run()
	execution := Execution{
		Output:    strings.ToValidUTF8(buffer.String(), ""),
		Bytes:     buffer.Len(),
		Truncated: buffer.Truncated(),
		ExitCode:  0,
	}
	if err == nil {
		return execution, nil
	}
	if ctxErr := ctx.Err(); ctxErr != nil {
		return execution, ctxErr
	}
	var exitError *exec.ExitError
	if errors.As(err, &exitError) {
		execution.ExitCode = exitError.ExitCode()
		return execution, nil
	}
	return execution, fmt.Errorf("start diagnostic executable: %w", err)
}

func isAbsoluteExecutablePath(value string) bool {
	return filepath.IsAbs(value) || path.IsAbs(value)
}

func cleanExecutablePath(value string) string {
	if path.IsAbs(value) && !filepath.IsAbs(value) {
		return path.Clean(value)
	}
	return filepath.Clean(value)
}

type limitedBuffer struct {
	mu        sync.Mutex
	content   []byte
	limit     int
	truncated bool
	reporter  func(OutputProgress)
}

func (b *limitedBuffer) Write(value []byte) (int, error) {
	b.mu.Lock()
	remaining := b.limit - len(b.content)
	if remaining > 0 {
		copyLength := min(remaining, len(value))
		b.content = append(b.content, value[:copyLength]...)
	}
	if len(value) > remaining {
		b.truncated = true
	}
	progress := OutputProgress{
		Output:    strings.ToValidUTF8(string(b.content), ""),
		Bytes:     len(b.content),
		Truncated: b.truncated,
	}
	reporter := b.reporter
	b.mu.Unlock()
	if reporter != nil {
		reporter(progress)
	}
	// Report the full write so os/exec can continue draining the child pipes
	// without retaining more output in memory.
	return len(value), nil
}

func (b *limitedBuffer) String() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	return string(b.content)
}

func (b *limitedBuffer) Len() int {
	b.mu.Lock()
	defer b.mu.Unlock()
	return len(b.content)
}

func (b *limitedBuffer) Truncated() bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.truncated
}