// ipgrep.go
package main

import (
	"bytes"
	"encoding/csv"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"sync"
	"time"
)

// IPInfo holds the data we want from ipinfo.io
type IPInfo struct {
	IP       string `json:"ip"`
	Hostname string `json:"hostname"`
	City     string `json:"city"`
	Region   string `json:"region"`
	Country  string `json:"country"`
	Org      string `json:"org"`
}

func main() {
	// Command-line flags.
	var (
		sortFlag    bool
		uniqFlag    bool
		macFlag     bool
		pingable    bool
		resolveFlag bool
		lookupFlag  bool
		fileName    string
	)
	flag.BoolVar(&uniqFlag, "u", false, "only show uniq IPs/MACs (implies -s)")
	flag.BoolVar(&sortFlag, "s", false, "sort output")
	flag.BoolVar(&macFlag, "m", false, "grep MAC-IDs instead of IPs")
	flag.BoolVar(&pingable, "p", false, "only show 'pingable' entries (MACs still beta)")
	flag.BoolVar(&resolveFlag, "r", false, "resolve (uses host for ip and arping for mac)")
	flag.BoolVar(&lookupFlag, "l", false, "lookup ip info using ipinfo.io and output CSV: ip,country,region,city,org,hostname")
	flag.StringVar(&fileName, "f", "", "input file")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [-u] [-s] [-m] [-p] [-r] [-l] [-f filename] [file...]\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()

	// If pingable is set, force sorting and uniqueness.
	if pingable || lookupFlag {
		sortFlag = true
		uniqFlag = true
	}

	if lookupFlag && macFlag {
		fmt.Fprintln(os.Stderr, "Lookup mode (-l) only works for IP addresses, not MAC addresses.")
		os.Exit(1)
	}

	// Regular expressions for IPs or MACs.
	var pattern string
	if macFlag {
		// Supports MAC formats: xx:xx:xx:xx:xx:xx or xxxx.xxxx.xxxx
		pattern = `(([a-fA-F0-9]{2}[:-]){5}[a-fA-F0-9]{2})|([a-fA-F0-9]{4}\.[a-fA-F0-9]{4}\.[a-fA-F0-9]{4})`
	} else {
		// Matches valid IPv4 addresses.
		pattern = `(((25[0-5])|(2[0-4][0-9])|([0-1]?\d?\d))\.){3}((25[0-5])|(2[0-4][0-9])|([0-1]?\d?\d))`
	}
	re, err := regexp.Compile(pattern)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error compiling regex: %v\n", err)
		os.Exit(1)
	}

	// Read input from -f file, extra args, or stdin.
	var inputData []byte
	if fileName != "" {
		inputData, err = ioutil.ReadFile(fileName)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error reading file %s: %v\n", fileName, err)
			os.Exit(1)
		}
	} else if flag.NArg() > 0 {
		var buf bytes.Buffer
		for _, fname := range flag.Args() {
			data, err := ioutil.ReadFile(fname)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error reading file %s: %v\n", fname, err)
				continue
			}
			buf.Write(data)
			buf.WriteByte('\n')
		}
		inputData = buf.Bytes()
	} else {
		inputData, err = io.ReadAll(os.Stdin)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error reading stdin: %v\n", err)
			os.Exit(1)
		}
	}

	// Filter matches using the regex.
	matches := re.FindAllString(string(inputData), -1)
	if matches == nil {
		os.Exit(0)
	}

	if sortFlag {
		sort.Strings(matches)
	}
	if uniqFlag {
		matches = unique(matches)
	}

	if pingable {
		matches = filterPingable(matches, macFlag)
		if sortFlag {
			sort.Strings(matches)
		}
	}

	// If lookup flag is set, perform ipinfo.io lookups with caching.
	if lookupFlag {
		cache := loadCache()
		var cacheMu sync.Mutex
		results := lookupIPInfo(matches, cache, &cacheMu)

		// Sort the results by IP.
		sort.Slice(results, func(i, j int) bool {
			return results[i].IP < results[j].IP
		})

		// Save the updated cache.
		saveCache(cache)

		// Output CSV using csv.Writer for proper CSV formatting.
		w := csv.NewWriter(os.Stdout)
		// Write header.
		if err := w.Write([]string{"ip", "country", "region", "city", "org", "hostname"}); err != nil {
			fmt.Fprintf(os.Stderr, "Error writing CSV header: %v\n", err)
			os.Exit(1)
		}
		for _, info := range results {
			record := []string{
				info.IP,
				info.Country,
				info.Region,
				info.City,
				info.Org,
				info.Hostname,
			}
			if err := w.Write(record); err != nil {
				fmt.Fprintf(os.Stderr, "Error writing CSV record: %v\n", err)
				os.Exit(1)
			}
		}
		w.Flush()
		if err := w.Error(); err != nil {
			fmt.Fprintf(os.Stderr, "Error flushing CSV data: %v\n", err)
		}
		return
	}

	// If resolve flag is set, perform resolution.
	if resolveFlag {
		results := resolveEntries(matches, macFlag)
		for _, r := range results {
			fmt.Print(r)
			if !macFlag {
				fmt.Println()
			}
		}
		return
	}

	// Otherwise, just output the matches.
	for _, m := range matches {
		fmt.Println(m)
	}
}

// unique removes duplicate strings from a slice.
func unique(input []string) []string {
	seen := make(map[string]struct{})
	var result []string
	for _, s := range input {
		if _, ok := seen[s]; !ok {
			seen[s] = struct{}{}
			result = append(result, s)
		}
	}
	return result
}

// filterPingable runs ping (or arping) concurrently on each entry.
func filterPingable(entries []string, mac bool) []string {
	var wg sync.WaitGroup
	var mu sync.Mutex
	var result []string

	for _, entry := range entries {
		wg.Add(1)
		go func(e string) {
			defer wg.Done()
			if isPingable(e, mac) {
				mu.Lock()
				result = append(result, e)
				mu.Unlock()
			}
		}(entry)
	}
	wg.Wait()
	return result
}

// isPingable tests if an entry is reachable using ping (or arping for MACs).
func isPingable(entry string, mac bool) bool {
	var cmd *exec.Cmd
	if mac {
		cmd = exec.Command("arping", "-c", "1", "-w", "5000000", entry)
	} else {
		cmd = exec.Command("ping", "-c", "1", "-w", "1", entry)
	}
	cmd.Stdout = nil
	cmd.Stderr = nil
	err := cmd.Run()
	return err == nil
}

// resolveEntries performs resolution via external commands.
func resolveEntries(entries []string, mac bool) []string {
	var wg sync.WaitGroup
	results := make([]string, len(entries))

	for i, entry := range entries {
		wg.Add(1)
		go func(i int, e string) {
			defer wg.Done()
			if mac {
				cmd := exec.Command("arping", "-q", "-c", "1", "-w", "5000000", e)
				if err := cmd.Run(); err == nil {
					cmd2 := exec.Command("arping", "-c", "1", e)
					out, err := cmd2.CombinedOutput()
					if err == nil {
						results[i] = string(out)
					} else {
						results[i] = e + "\n"
					}
				} else {
					results[i] = e + "\n"
				}
			} else {
				cmd := exec.Command("host", e)
				out, err := cmd.CombinedOutput()
				if err == nil {
					// Extract the hostname via regex (similar to grep -Po '(?<=pointer ).*')
					reHost := regexp.MustCompile(`(?i)(?<=pointer\s)(\S+)`)
					match := reHost.FindString(string(out))
					if match != "" {
						results[i] = fmt.Sprintf("%s  %s", e, match)
					} else {
						results[i] = e
					}
				} else {
					results[i] = e
				}
			}
		}(i, entry)
	}
	wg.Wait()
	return results
}

// lookupIPInfo queries ipinfo.io for each IP concurrently,
// checking a local cache before going to the network.
// It returns a slice of IPInfo.
func lookupIPInfo(entries []string, cache map[string]IPInfo, cacheMu *sync.Mutex) []IPInfo {
	var wg sync.WaitGroup
	var mu sync.Mutex
	var results []IPInfo

	client := &http.Client{
		Timeout: 5 * time.Second,
	}

	for _, ip := range entries {
		wg.Add(1)
		go func(ip string) {
			defer wg.Done()
			// Check cache first.
			cacheMu.Lock()
			info, found := cache[ip]
			cacheMu.Unlock()
			if found {
				mu.Lock()
				results = append(results, info)
				mu.Unlock()
				return
			}

			// Not in cache; perform HTTP lookup.
			url := fmt.Sprintf("https://ipinfo.io/%s", ip)
			resp, err := client.Get(url)
			if err != nil {
				return
			}
			defer resp.Body.Close()
			body, err := io.ReadAll(resp.Body)
			if err != nil {
				return
			}
			var newInfo IPInfo
			if err := json.Unmarshal(body, &newInfo); err != nil {
				return
			}
			// Only add valid responses.
			if newInfo.IP == "" {
				return
			}
			// Update cache.
			cacheMu.Lock()
			cache[ip] = newInfo
			cacheMu.Unlock()
			mu.Lock()
			results = append(results, newInfo)
			mu.Unlock()
		}(ip)
	}
	wg.Wait()
	return results
}

// loadCache reads the cache from ~/.ipgrep.db (if present)
// and returns it as a map[string]IPInfo.
func loadCache() map[string]IPInfo {
	home, err := os.UserHomeDir()
	if err != nil {
		return make(map[string]IPInfo)
	}
	cachePath := home + "/.ipgrep.db"
	data, err := ioutil.ReadFile(cachePath)
	if err != nil {
		// File doesn't exist or can't be read, start with an empty cache.
		return make(map[string]IPInfo)
	}
	var cache map[string]IPInfo
	if err := json.Unmarshal(data, &cache); err != nil {
		// If unmarshal fails, use an empty cache.
		return make(map[string]IPInfo)
	}
	return cache
}

// saveCache writes the cache map to ~/.ipgrep.db.
func saveCache(cache map[string]IPInfo) {
	home, err := os.UserHomeDir()
	if err != nil {
		return
	}
	cachePath := home + "/.ipgrep.db"
	data, err := json.MarshalIndent(cache, "", "  ")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error marshaling cache: %v\n", err)
		return
	}
	if err := ioutil.WriteFile(cachePath, data, 0644); err != nil {
		fmt.Fprintf(os.Stderr, "Error writing cache file: %v\n", err)
	}
}