package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/go-ping/ping"
	"gopkg.in/yaml.v3"
)

type Config struct {
	InputFile       string `yaml:"input_file"`
	OutputFile      string `yaml:"output_file"`
	IntervalSeconds int    `yaml:"interval_seconds"`
	CooldownMinutes int    `yaml:"cooldown_minutes"`
	EnableTraceroute bool  `yaml:"enable_traceroute"`
	TracerouteMaxHops int  `yaml:"traceroute_max_hops"`
}

type PingResult struct {
	IP         string        `json:"ip"`
	Sent       int           `json:"sent"`
	Received   int           `json:"received"`
	PacketLoss float64       `json:"packet_loss"`
	AvgRtt     time.Duration `json:"avg_rtt"`
	Timestamp  time.Time     `json:"timestamp"`
	Error      string        `json:"error,omitempty"`
	Traceroute *TracerouteResult `json:"traceroute,omitempty"`
}

type TracerouteResult struct {
	Method    string         `json:"method"` // "icmp" or "tcp"
	Hops      []TracerouteHop `json:"hops"`
	Completed bool           `json:"completed"`
	Error     string         `json:"error,omitempty"`
}

type TracerouteHop struct {
	TTL     int           `json:"ttl"`
	IP      string        `json:"ip"`
	Rtt     time.Duration `json:"rtt,omitempty"`
	Timeout bool          `json:"timeout,omitempty"`
}

var (
	// cooldownCache stores IP -> last ping time
	cooldownCache = make(map[string]time.Time)
	cacheMux      sync.Mutex
	verbose       bool
)

func main() {
	// CLI flags
	configPath := flag.String("config", "config.yaml", "Path to config file")
	verboseFlag := flag.Bool("v", false, "Enable verbose logging")
	flag.BoolVar(verboseFlag, "verbose", false, "Enable verbose logging")
	help := flag.Bool("help", false, "Show help message")
	flag.Parse()

	if *help {
		fmt.Println("Ping Service - Monitor network endpoints via ping")
		fmt.Println("\nUsage:")
		flag.PrintDefaults()
		fmt.Println("\nConfig file format (YAML):")
		fmt.Println("  input_file: Path/URL to file containing IPs (one per line)")
		fmt.Println("  output_file: Path/URL/socket for JSON results")
		fmt.Println("  interval_seconds: How often to check for new IPs")
		fmt.Println("  cooldown_minutes: Minimum time between pings for same IP")
		fmt.Println("  enable_traceroute: Enable traceroute after successful ping")
		fmt.Println("  traceroute_max_hops: Maximum TTL for traceroute (default: 30)")
		os.Exit(0)
	}

	verbose = *verboseFlag

	config := loadConfig(*configPath)
	if config.CooldownMinutes == 0 {
		config.CooldownMinutes = 10
	}
	if config.TracerouteMaxHops == 0 {
		config.TracerouteMaxHops = 30
	}

	ticker := time.NewTicker(time.Duration(config.IntervalSeconds) * time.Second)
	log.Printf("App started. Cooldown set to %d mins. Polling every %ds...", config.CooldownMinutes, config.IntervalSeconds)
	if config.EnableTraceroute {
		log.Printf("Traceroute enabled (max %d hops)", config.TracerouteMaxHops)
	}

	for {
		process(config)
		<-ticker.C
	}
}

func loadConfig(path string) *Config {
	f, err := os.ReadFile(path)
	if err != nil {
		log.Fatalf("Error reading config: %v", err)
	}
	var cfg Config
	err = yaml.Unmarshal(f, &cfg)
	if err != nil {
		log.Fatalf("Error parsing config: %v", err)
	}
	return &cfg
}

func process(cfg *Config) {
	// 1. Read IPs
	data, err := readSource(cfg.InputFile)
	if err != nil {
		log.Printf("Input Error: %v", err)
		return
	}
	rawIps := strings.Fields(string(data))

	if verbose {
		log.Printf("Read %d IPs from %s", len(rawIps), cfg.InputFile)
	}

	// Filter IPs based on cooldown
	var ipsToPing []string
	cacheMux.Lock()
	for _, ip := range rawIps {
		if last, ok := cooldownCache[ip]; ok {
			if time.Since(last) < time.Duration(cfg.CooldownMinutes)*time.Minute {
				if verbose {
					log.Printf("Skipping %s (cooldown: %v remaining)", ip, time.Duration(cfg.CooldownMinutes)*time.Minute-time.Since(last))
				}
				continue // Skip this IP
			}
		}
		ipsToPing = append(ipsToPing, ip)
		cooldownCache[ip] = time.Now()
	}
	cacheMux.Unlock()

	if len(ipsToPing) == 0 {
		if verbose {
			log.Println("No IPs to ping (all in cooldown)")
		}
		return
	}

	if verbose {
		log.Printf("Pinging %d IPs: %v", len(ipsToPing), ipsToPing)
	}

	// 2. Perform Pings
	var wg sync.WaitGroup
	results := make([]PingResult, len(ipsToPing))

	for i, ip := range ipsToPing {
		wg.Add(1)
		go func(idx int, targetIP string) {
			defer wg.Done()
			results[idx] = runPing(targetIP)

			// If ping successful and traceroute enabled, do traceroute
			if cfg.EnableTraceroute && results[idx].Error == "" && results[idx].Received > 0 {
				if verbose {
					log.Printf("Running traceroute to %s...", targetIP)
				}
				results[idx].Traceroute = runTraceroute(targetIP, cfg.TracerouteMaxHops)
			}

			if verbose {
				if results[idx].Error != "" {
					log.Printf("Pinged %s: ERROR - %s", results[idx].IP, results[idx].Error)
				} else {
					log.Printf("Pinged %s: %d/%d packets, %.2f%% loss, avg RTT: %v",
						results[idx].IP, results[idx].Received, results[idx].Sent,
						results[idx].PacketLoss, results[idx].AvgRtt)
					if results[idx].Traceroute != nil {
						log.Printf("  Traceroute: %d hops via %s", len(results[idx].Traceroute.Hops), results[idx].Traceroute.Method)
					}
				}
			}
		}(i, ip)
	}
	wg.Wait()

	// 3. Write Output
	outputData, _ := json.MarshalIndent(results, "", "  ")
	err = writeDestination(cfg.OutputFile, outputData)
	if err != nil {
		log.Printf("Output Error: %v", err)
	} else if verbose {
		log.Printf("Wrote results to %s", cfg.OutputFile)
	}
}

// Logic: If socket exists, dial it. If not, listen on it.
func handleSocket(path string, data []byte, mode string) ([]byte, error) {
	if _, err := os.Stat(path); err == nil {
		// SOCKET EXISTS: Connect as Client
		conn, err := net.DialTimeout("unix", path, 2*time.Second)
		if err != nil {
			// If we can't connect, the socket might be "stale" (file exists but no one listening)
			os.Remove(path)
			return handleSocket(path, data, mode) // Recursive call to create listener
		}
		defer conn.Close()

		if mode == "read" {
			return io.ReadAll(conn)
		} else {
			_, err = conn.Write(data)
			return nil, err
		}
	} else {
		// SOCKET DOES NOT EXIST: Create as Server
		l, err := net.Listen("unix", path)
		if err != nil {
			return nil, err
		}
		defer l.Close()
		defer os.Remove(path)

		// Set a timeout so the app doesn't hang forever if no one connects
		l.(*net.UnixListener).SetDeadline(time.Now().Add(10 * time.Second))

		conn, err := l.Accept()
		if err != nil {
			return nil, fmt.Errorf("timeout waiting for socket connection: %v", err)
		}
		defer conn.Close()

		if mode == "read" {
			return io.ReadAll(conn)
		} else {
			_, err = conn.Write(data)
			return nil, err
		}
	}
}

func readSource(src string) ([]byte, error) {
	if strings.HasPrefix(src, "http") {
		resp, err := http.Get(src)
		if err != nil {
			return nil, err
		}
		defer resp.Body.Close()
		return io.ReadAll(resp.Body)
	}

	if strings.HasSuffix(src, ".sock") {
		return handleSocket(src, nil, "read")
	}

	return os.ReadFile(src)
}

func writeDestination(dest string, data []byte) error {
	if strings.HasPrefix(dest, "http") {
		resp, err := http.Post(dest, "application/json", bytes.NewBuffer(data))
		if err != nil {
			return err
		}
		defer resp.Body.Close()
		return nil
	}

	if strings.HasSuffix(dest, ".sock") {
		_, err := handleSocket(dest, data, "write")
		return err
	}

	return os.WriteFile(dest, data, 0644)
}

func runPing(ip string) PingResult {
	pinger, err := ping.NewPinger(ip)
	res := PingResult{IP: ip, Timestamp: time.Now()}
	if err != nil {
		res.Error = err.Error()
		return res
	}

	pinger.Count = 3
	pinger.Timeout = time.Second * 5
	// pinger.SetPrivileged(true) // Uncomment if running on Linux/Windows as admin

	err = pinger.Run()
	if err != nil {
		res.Error = err.Error()
		return res
	}

	stats := pinger.Statistics()
	res.Sent = stats.PacketsSent
	res.Received = stats.PacketsRecv
	res.PacketLoss = stats.PacketLoss
	res.AvgRtt = stats.AvgRtt
	return res
}

func runTraceroute(ip string, maxHops int) *TracerouteResult {
	result := &TracerouteResult{
		Method: "icmp",
		Hops:   []TracerouteHop{},
	}

	// Try ICMP traceroute first
	if tr := tryICMPTraceroute(ip, maxHops); tr != nil {
		return tr
	}

	// If ICMP fails, try TCP traceroute on common ports
	if verbose {
		log.Printf("ICMP traceroute failed for %s, trying TCP...", ip)
	}

	for _, port := range []int{80, 443, 22} {
		if tr := tryTCPTraceroute(ip, port, maxHops); tr != nil {
			tr.Method = fmt.Sprintf("tcp/%d", port)
			return tr
		}
	}

	result.Error = "traceroute failed (both ICMP and TCP)"
	return result
}

func tryICMPTraceroute(ip string, maxHops int) *TracerouteResult {
	var cmd *exec.Cmd

	// Detect OS and use appropriate command
	switch {
	case fileExists("/usr/bin/traceroute"):
		cmd = exec.Command("traceroute", "-m", strconv.Itoa(maxHops), "-n", "-q", "1", ip)
	case fileExists("/usr/sbin/traceroute"):
		cmd = exec.Command("traceroute", "-m", strconv.Itoa(maxHops), "-n", "-q", "1", ip)
	case fileExists("/bin/traceroute"):
		cmd = exec.Command("traceroute", "-m", strconv.Itoa(maxHops), "-n", "-q", "1", ip)
	default:
		// Try without full path
		cmd = exec.Command("traceroute", "-m", strconv.Itoa(maxHops), "-n", "-q", "1", ip)
	}

	output, err := cmd.CombinedOutput()
	if err != nil {
		return nil
	}

	return parseTracerouteOutput(string(output), "icmp")
}

func tryTCPTraceroute(ip string, port int, maxHops int) *TracerouteResult {
	var cmd *exec.Cmd

	// Try tcptraceroute if available
	if fileExists("/usr/bin/tcptraceroute") || fileExists("/usr/sbin/tcptraceroute") {
		cmd = exec.Command("tcptraceroute", "-m", strconv.Itoa(maxHops), "-n", "-q", "1", ip, strconv.Itoa(port))
	} else {
		// Try traceroute with TCP (-T flag on Linux)
		cmd = exec.Command("traceroute", "-T", "-p", strconv.Itoa(port), "-m", strconv.Itoa(maxHops), "-n", "-q", "1", ip)
	}

	output, err := cmd.CombinedOutput()
	if err != nil {
		return nil
	}

	return parseTracerouteOutput(string(output), fmt.Sprintf("tcp/%d", port))
}

func parseTracerouteOutput(output string, method string) *TracerouteResult {
	result := &TracerouteResult{
		Method: method,
		Hops:   []TracerouteHop{},
		Completed: false,
	}

	lines := strings.Split(output, "\n")

	// Regex to match: " 1  192.168.1.1  1.234 ms"
	hopRegex := regexp.MustCompile(`^\s*(\d+)\s+([0-9.]+)\s+([0-9.]+)\s*ms`)
	// Regex to match timeouts: " 1  * * *"
	timeoutRegex := regexp.MustCompile(`^\s*(\d+)\s+\*`)

	for _, line := range lines {
		if match := hopRegex.FindStringSubmatch(line); match != nil {
			ttl, _ := strconv.Atoi(match[1])
			hopIP := match[2]
			rttMs, _ := strconv.ParseFloat(match[3], 64)
			rtt := time.Duration(rttMs * float64(time.Millisecond))

			result.Hops = append(result.Hops, TracerouteHop{
				TTL: ttl,
				IP:  hopIP,
				Rtt: rtt,
			})
		} else if match := timeoutRegex.FindStringSubmatch(line); match != nil {
			ttl, _ := strconv.Atoi(match[1])
			result.Hops = append(result.Hops, TracerouteHop{
				TTL:     ttl,
				Timeout: true,
			})
		}
	}

	if len(result.Hops) > 0 {
		result.Completed = true
	}

	return result
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}