package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path"
	"strings"
)

type IPAddress struct {
	Address string `json:"address"`
	DNSName string `json:"dns_name"`
}

type NetBoxResponse struct {
	Results []IPAddress `json:"results"`
}

func FetchNetboxIPAddresses(apiBaseURL, token string) ([]IPAddress, error) {
	u, _ := url.Parse(apiBaseURL)
	u.Path = path.Join(u.Path, "ipam/ip-addresses")
	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Authorization", "Token "+token)

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("unexpected status: %s, body: %s", resp.Status, string(body))
	}

	var nbResp NetBoxResponse
	decoder := json.NewDecoder(resp.Body)
	if err := decoder.Decode(&nbResp); err != nil {
		return nil, err
	}

	// Filter out entries with empty DNSName and strip CIDR from Address
	var filtered []IPAddress
	for _, ip := range nbResp.Results {
		if ip.DNSName != "" {
			ip.Address = strings.Split(ip.Address, "/")[0]
			filtered = append(filtered, ip)
		}
	}
	return filtered, nil
}

var apiBaseURL string

func main() {
	if len(os.Args) < 2 {
		fmt.Fprintf(os.Stderr, "Usage: %s <run|serve> [flags]\n", os.Args[0])
		os.Exit(1)
	}

	subCmd := os.Args[1]
	// Prepare a FlagSet for subcommands, sharing the same flags
	flags := flag.NewFlagSet(subCmd, flag.ExitOnError)

	var apiURL string
	var outputFormat string
	var outputFile string
	var outputDir string
	var dryRun bool
	var listenAddr string

	flags.StringVar(&apiURL, "api-url", "https://netbox.koszewscy.waw.pl/api", "NetBox API URL to fetch IP addresses")
	flags.StringVar(&outputFormat, "output-format", "hosts", "output format: config or hosts")
	flags.StringVar(&outputFile, "output-file", "", "output file name (overrides default)")
	flags.StringVar(&outputDir, "output-dir", "", "output base directory (overrides default)")
	flags.BoolVar(&dryRun, "dry-run", false, "if set, do not write to dnsmasq config file, just print the output")
	flags.StringVar(&listenAddr, "listen", ":8080", "address and port to listen on (e.g. :8080 or 127.0.0.1:8080)")

	flags.Parse(os.Args[2:])
	apiBaseURL = apiURL

	// Set defaults based on output format
	var defaultDir, defaultFile string
	switch outputFormat {
	case "hosts":
		defaultDir = "/etc"
		defaultFile = "hosts.netbox"
	case "config":
		defaultDir = "/etc/dnsmasq.d"
		defaultFile = "netbox.conf"
	default:
		log.Fatalf("Unknown output format: %s", outputFormat)
	}
	dir := defaultDir
	file := defaultFile
	if outputDir != "" {
		dir = outputDir
	}
	if outputFile != "" {
		file = outputFile
	}

	switch subCmd {
	case "run":
		if listenAddr != ":8080" && listenAddr != "" {
			log.Fatalf("The --listen flag is only valid with the 'serve' subcommand.")
		}
		if dryRun {
			log.Println("Dry run mode enabled, not writing to dnsmasq config file.")
			CreateDnsMasqConfig(true, outputFormat, dir, file)
		} else {
			CreateDnsMasqConfig(false, outputFormat, dir, file)
		}
	case "serve":
		http.HandleFunc("/update-dnsmasq", func(w http.ResponseWriter, r *http.Request) {
			if r.Method != http.MethodGet {
				w.WriteHeader(http.StatusMethodNotAllowed)
				w.Write([]byte("Method not allowed"))
				return
			}
			CreateDnsMasqConfig(false, outputFormat, dir, file)
			cmd := exec.Command("systemctl", "restart", "dnsmasq.service")
			if err := cmd.Run(); err != nil {
				log.Printf("Failed to restart dnsmasq.service: %v", err)
				w.Write([]byte("\nWarning: Failed to restart dnsmasq.service.\n"))
			}
			w.Write([]byte("DNSMasq config updated."))
		})
		log.Printf("Starting web service on %s...", listenAddr)
		if err := http.ListenAndServe(listenAddr, nil); err != nil {
			log.Fatalf("Failed to start server: %v", err)
		}
	default:
		fmt.Fprintf(os.Stderr, "Unknown subcommand: %s\n", subCmd)
		os.Exit(1)
	}
}

// Update CreateDnsMasqConfig to accept new parameters
func CreateDnsMasqConfig(writeConfig bool, outputFormat, dir, file string) {
	token := os.Getenv("NETBOX_TOKEN")
	if token == "" {
		home := os.Getenv("HOME")
		paths := []string{
			home + "/.netbox/token",
			"/etc/netbox/token",
		}
		for _, path := range paths {
			data, err := os.ReadFile(path)
			if err == nil {
				token = strings.TrimSpace(string(data))
				break
			}
		}
		if token == "" {
			log.Fatal("NETBOX_TOKEN not set and no token file found")
		}
	}

	ips, err := FetchNetboxIPAddresses(apiBaseURL, token)
	if err != nil {
		log.Fatalf("Error fetching IP addresses: %v", err)
	}

	if stat, err := os.Stat(dir); err == nil && stat.IsDir() && !writeConfig {
		f, err := os.Create(path.Join(dir, file))
		if err != nil {
			log.Fatalf("Failed to create %s: %v", file, err)
		}
		defer f.Close()
		for _, ip := range ips {
			if outputFormat == "hosts" {
				fmt.Fprintf(f, "%s\t%s\n", ip.Address, ip.DNSName)
			} else {
				fmt.Fprintf(f, "address=/%s/%s\n", ip.DNSName, ip.Address)
			}
		}
	} else {
		for _, ip := range ips {
			if outputFormat == "hosts" {
				fmt.Printf("%s\t%s\n", ip.Address, ip.DNSName)
			} else {
				fmt.Printf("address=/%s/%s\n", ip.DNSName, ip.Address)
			}
		}
	}
}