package main

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"
)

type OrgResponse struct {
	Orgname string `json:"orgname"`
	Source  string `json:"source"`
	Cached  bool   `json:"cached"`
}

type PrefixResponse struct {
	Prefixes []struct {
		Prefix string `json:"Prefix"`
		Count  int    `json:"Count"`
		Total  int    `json:"Total"`
	} `json:"prefixes"`
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: go run main.go [--ipv4-only|--ipv6-only] [--minify] <AS_NUMBER1> [AS_NUMBER2 ...]")
		fmt.Println("Options:")
		fmt.Println("  --ipv4-only    Only output IPv4 prefixes")
		fmt.Println("  --ipv6-only    Only output IPv6 prefixes")
		fmt.Println("  --minify       Output minified format to ip-asn-minified.txt")
		os.Exit(1)
	}

	// Parse flags
	var ipv4Only, ipv6Only, minify bool
	var asNumbers []string

	// Process arguments
	for _, arg := range os.Args[1:] {
		switch arg {
		case "--ipv4-only":
			ipv4Only = true
		case "--ipv6-only":
			ipv6Only = true
		case "--minify":
			minify = true
		default:
			asNumbers = append(asNumbers, arg)
		}
	}

	if ipv4Only && ipv6Only {
		fmt.Println("Error: Cannot use both --ipv4-only and --ipv6-only at the same time")
		os.Exit(1)
	}

	if len(asNumbers) == 0 {
		fmt.Println("Error: No AS numbers provided")
		os.Exit(1)
	}

	// Create output file
	outputFile, err := os.Create("ip-asn.txt")
	if err != nil {
		fmt.Printf("Error creating output file: %v\n", err)
		return
	}
	defer outputFile.Close()

	// Create minified output file if needed
	var minifiedFile *os.File
	if minify {
		minifiedFile, err = os.Create("ip-asn-minified.txt")
		if err != nil {
			fmt.Printf("Error creating minified output file: %v\n", err)
			return
		}
		defer minifiedFile.Close()
	}

	// Create HTTP client
	client := &http.Client{}

	// Process each AS number
	for i, arg := range asNumbers {
		asNumber := strings.ToLower(strings.TrimPrefix(strings.TrimPrefix(arg, "AS"), "as"))
		fmt.Printf("Processing AS%s (%d/%d)...\n", asNumber, i+1, len(asNumbers))

		// Random wait between requests (1-3 seconds)
		if i > 0 {
			waitTime := time.Duration(1000+rand.Intn(2000)) * time.Millisecond
			fmt.Printf("Waiting for %v before next request...\n", waitTime)
			time.Sleep(waitTime)
		}

		// Write AS separator and header
		if i > 0 {
			fmt.Fprintf(outputFile, "\n")
		}

		// Get organization name
		fmt.Printf("Fetching organization info for AS%s...\n", asNumber)
		orgReq, err := http.NewRequest("GET", fmt.Sprintf("https://bgp.he.net/orgname/as%s", asNumber), nil)
		if err != nil {
			fmt.Printf("Error creating request: %v\n", err)
			continue
		}

		// Set required headers
		orgReq.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:133.0) Gecko/20100101 Firefox/133.0")
		orgReq.Header.Set("Accept", "*/*")
		orgReq.Header.Set("Accept-Language", "en-US,en;q=0.5")
		orgReq.Header.Set("Accept-Encoding", "gzip")
		orgReq.Header.Set("Referer", fmt.Sprintf("https://bgp.he.net/AS%s", asNumber))

		orgResp, err := client.Do(orgReq)
		if err != nil {
			fmt.Printf("Error making request: %v\n", err)
			continue
		}
		defer orgResp.Body.Close()

		orgData, err := readResponse(orgResp)
		if err != nil {
			fmt.Printf("Error reading response: %v\n", err)
			continue
		}

		var orgInfo OrgResponse
		if err := json.Unmarshal(orgData, &orgInfo); err != nil {
			fmt.Printf("Error parsing JSON: %v\n", err)
			continue
		}

		fmt.Fprintf(outputFile, "# =====================================\n")
		fmt.Fprintf(outputFile, "# AS%s Prefixes\n", asNumber)
		fmt.Fprintf(outputFile, "# Organization: %s\n", orgInfo.Orgname)
		fmt.Fprintf(outputFile, "# =====================================\n\n")

		// Get prefixes
		fmt.Printf("Fetching prefixes for AS%s...\n", asNumber)
		prefixReq, err := http.NewRequest("GET", fmt.Sprintf("https://bgp.he.net/super-lg/report/api/v1/prefixes/originated/%s", asNumber), nil)
		if err != nil {
			fmt.Printf("Error creating request: %v\n", err)
			continue
		}

		// Set the same headers for prefix request
		prefixReq.Header = orgReq.Header

		prefixResp, err := client.Do(prefixReq)
		if err != nil {
			fmt.Printf("Error making request: %v\n", err)
			continue
		}
		defer prefixResp.Body.Close()

		prefixData, err := readResponse(prefixResp)
		if err != nil {
			fmt.Printf("Error reading response: %v\n", err)
			continue
		}

		var prefixInfo PrefixResponse
		if err := json.Unmarshal(prefixData, &prefixInfo); err != nil {
			fmt.Printf("Error parsing JSON: %v\n", err)
			continue
		}

		// Separate IPv4 and IPv6 prefixes
		var ipv4Prefixes []string
		var ipv6Prefixes []string

		for _, p := range prefixInfo.Prefixes {
			if strings.Contains(p.Prefix, ":") {
				ipv6Prefixes = append(ipv6Prefixes, p.Prefix)
			} else {
				ipv4Prefixes = append(ipv4Prefixes, p.Prefix)
			}
		}

		// Sort prefixes
		sort.Strings(ipv4Prefixes)
		sort.Strings(ipv6Prefixes)

		// Write to minified file if enabled
		if minify {
			if i > 0 {
				fmt.Fprint(minifiedFile, " ")
			}
			if !ipv6Only {
				for i, prefix := range ipv4Prefixes {
					if i > 0 {
						fmt.Fprint(minifiedFile, " ")
					}
					fmt.Fprint(minifiedFile, prefix)
				}
			}
			if !ipv4Only && len(ipv6Prefixes) > 0 {
				if len(ipv4Prefixes) > 0 && !ipv6Only {
					fmt.Fprint(minifiedFile, " ")
				}
				for i, prefix := range ipv6Prefixes {
					if i > 0 {
						fmt.Fprint(minifiedFile, " ")
					}
					fmt.Fprint(minifiedFile, prefix)
				}
			}
		}

		// Write IPv4 prefixes
		if len(ipv4Prefixes) > 0 && !ipv6Only {
			fmt.Fprintf(outputFile, "## AS%s - %s - IPv4 Prefixes\n", asNumber, orgInfo.Orgname)
			for _, prefix := range ipv4Prefixes {
				fmt.Fprintln(outputFile, prefix)
			}
		}

		// Write IPv6 prefixes
		if len(ipv6Prefixes) > 0 && !ipv4Only {
			if len(ipv4Prefixes) > 0 && !ipv6Only {
				fmt.Fprintln(outputFile, "")
			}
			fmt.Fprintf(outputFile, "## AS%s - %s - IPv6 Prefixes\n", asNumber, orgInfo.Orgname)
			for _, prefix := range ipv6Prefixes {
				fmt.Fprintln(outputFile, prefix)
			}
		}

		fmt.Printf("Completed processing AS%s\n", asNumber)
	}

	fmt.Println("All AS numbers processed successfully!")
}

func readResponse(resp *http.Response) ([]byte, error) {
	var reader io.ReadCloser = resp.Body
	if resp.Header.Get("Content-Encoding") == "gzip" {
		var err error
		reader, err = gzip.NewReader(resp.Body)
		if err != nil {
			return nil, fmt.Errorf("error creating gzip reader: %v", err)
		}
		defer reader.Close()
	}
	return io.ReadAll(reader)
}