package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"math"
	"os"
	"sync"
	"time"

	"github.com/jsdelivr/globalping-cli/globalping"
	"github.com/knadh/koanf/providers/posflag"
	"github.com/knadh/koanf/v2"
	"github.com/mr-karan/doggo/internal/app"
	"github.com/mr-karan/doggo/pkg/resolvers"
	"github.com/mr-karan/doggo/pkg/utils"
	flag "github.com/spf13/pflag"
)

// Exit codes used by the CLI. Exit 0 is implicit success; partial success
// (some resolvers answered, others failed) is 2; full lookup failure remains
// 9 to preserve compatibility with the pre-existing convention.
const (
	exitGenericFailure = 1
	exitPartialFailure = 2
	exitLookupFailure  = 9
)

var (
	buildVersion = "unknown"
	buildDate    = "unknown"
	k            = koanf.New(".")
)

func main() {
	if len(os.Args) > 1 && os.Args[1] == "completions" {
		completionsCommand()
		return
	}

	cfg, err := loadConfig()
	if err != nil {
		fmt.Printf("Error loading configuration: %v\n", err)
		os.Exit(1)
	}

	if cfg.showVersion {
		fmt.Printf("%s - %s\n", buildVersion, buildDate)
		os.Exit(0)
	}

	logger := utils.InitLogger(cfg.debug)
	app := initializeApp(logger, cfg)

	if app.QueryFlags.GPFrom != "" {
		res, err := app.GlobalpingMeasurement()
		if err != nil {
			logger.Error("Error fetching globalping measurement", "error", err)
			os.Exit(2)
		}
		if app.QueryFlags.ShowJSON {
			err = app.OutputGlobalpingJSON(res)
		} else if app.QueryFlags.ShortOutput {
			err = app.OutputGlobalpingShort(res)
		} else {
			err = app.OutputGlobalping(res)
		}
		if err != nil {
			logger.Error("Error outputting globalping measurement", "error", err)
			os.Exit(2)
		}
		return
	}

	if cfg.reverseLookup {
		app.ReverseLookup()
	}

	app.LoadFallbacks()
	app.PrepareQuestions()

	if err := app.LoadNameservers(); err != nil {
		logger.Error("Error loading nameservers", "error", err)
		os.Exit(2)
	}

	resolvers, err := loadResolvers(app, cfg)
	if err != nil {
		logger.Error("Error loading resolvers", "error", err)
		os.Exit(2)
	}
	app.Resolvers = resolvers

	if len(app.QueryFlags.QNames) == 0 {
		cfg.flagSet.Usage()
		os.Exit(0)
	}

	responses, lookupErrors := performLookup(app, cfg)
	outputResults(app, responses, lookupErrors)
}

type config struct {
	flagSet       *flag.FlagSet
	showVersion   bool
	debug         bool
	reverseLookup bool
	timeout       time.Duration
	queryFlags    resolvers.QueryFlags
	outputJSON    bool
	showTime      bool
	useColor      bool
}

func loadConfig() (*config, error) {
	cfg := &config{}
	cfg.flagSet = setupFlags()

	if err := parseAndLoadFlags(cfg.flagSet); err != nil {
		return nil, fmt.Errorf("error parsing or loading flags: %w", err)
	}

	cfg.showVersion = k.Bool("version")
	cfg.debug = k.Bool("debug")
	cfg.reverseLookup = k.Bool("reverse")
	cfg.timeout = k.Duration("timeout")
	cfg.outputJSON = k.Bool("json")
	cfg.showTime = k.Bool("time")
	cfg.useColor = k.Bool("color")

	bufsize := k.Int("bufsize")
	if bufsize < 0 || bufsize > 65535 {
		return nil, fmt.Errorf("--bufsize must be between 0 and 65535, got %d", bufsize)
	}
	if bufsize > 0 && bufsize < 512 {
		return nil, fmt.Errorf("--bufsize must be 0 or at least 512 (RFC 6891), got %d", bufsize)
	}

	cfg.queryFlags = resolvers.QueryFlags{
		AA: k.Bool("aa"),
		AD: k.Bool("ad"),
		CD: k.Bool("cd"),
		RD: k.Bool("rd"),
		Z:  k.Bool("z"),
		DO: k.Bool("do"),

		// EDNS0 options
		NSID:    k.Bool("nsid"),
		Cookie:  k.Bool("cookie"),
		Padding: k.Bool("padding"),
		EDE:     k.Bool("ede"),
		ECS:     k.String("ecs"),
		Bufsize: uint16(bufsize),
	}

	return cfg, nil
}

func setupFlags() *flag.FlagSet {
	f := flag.NewFlagSet("config", flag.ContinueOnError)
	f.Usage = renderCustomHelp

	f.StringSliceP("query", "q", []string{}, "Domain name to query")
	f.StringSliceP("type", "t", []string{}, "Type of DNS record to be queried (A, AAAA, MX etc)")
	f.StringSliceP("class", "c", []string{}, "Network class of the DNS record to be queried (IN, CH, HS etc)")
	f.StringSliceP("nameserver", "n", []string{}, "Address of the nameserver to send packets to")
	f.BoolP("reverse", "x", false, "Performs a DNS Lookup for an IPv4 or IPv6 address")

	f.String("gp-from", "", "Probe locations as a comma-separated list")
	f.Int("gp-limit", 1, "Limit the number of probes to use")

	f.DurationP("timeout", "T", 5*time.Second, "Sets the timeout for a query")
	f.Bool("search", true, "Use the search list provided in resolv.conf")
	f.Int("ndots", -1, "Specify the ndots parameter")
	f.BoolP("ipv4", "4", false, "Use IPv4 only")
	f.BoolP("ipv6", "6", false, "Use IPv6 only")
	f.String("strategy", "all", "Strategy to query nameservers (all, random, first, internal)")
	f.String("tls-hostname", "", "Hostname for certificate verification")
	f.Bool("skip-hostname-verification", false, "Skip TLS Hostname Verification")

	f.Bool("any", false, "Query all supported DNS record types")

	f.BoolP("json", "J", false, "Set the output format as JSON")
	f.Bool("short", false, "Short output format")
	f.Bool("time", false, "Display how long the response took")
	f.Bool("color", true, "Show colored output")
	f.Bool("debug", false, "Enable debug mode")

	// Add flags for DNS query options
	f.Bool("aa", false, "Set Authoritative Answer flag")
	f.Bool("ad", false, "Set Authenticated Data flag")
	f.Bool("cd", false, "Set Checking Disabled flag")
	f.Bool("rd", true, "Set Recursion Desired flag (default: true)")
	f.Bool("z", false, "Set Z flag (reserved for future use)")
	f.Bool("do", false, "Set DNSSEC OK flag")

	// Add flags for EDNS0 options
	f.Bool("nsid", false, "Request Name Server Identifier (NSID)")
	f.Bool("cookie", false, "Request DNS Cookie")
	f.Bool("padding", false, "Request EDNS padding for privacy")
	f.Bool("ede", false, "Request Extended DNS Errors")
	f.String("ecs", "", "EDNS Client Subnet (e.g., '192.0.2.0/24' or '2001:db8::/32')")
	f.Int("bufsize", 0, "EDNS UDP buffer size in bytes (512-65535); setting this enables EDNS even without other EDNS options. Default is 1232 when EDNS is enabled.")

	f.Bool("version", false, "Show version of doggo")

	return f
}

func parseAndLoadFlags(f *flag.FlagSet) error {
	if err := f.Parse(os.Args[1:]); err != nil {
		return fmt.Errorf("error parsing flags: %w", err)
	}
	if err := k.Load(posflag.Provider(f, ".", k), nil); err != nil {
		return fmt.Errorf("error loading flags: %w", err)
	}
	return nil
}

func initializeApp(logger *slog.Logger, cfg *config) *app.App {
	gpConfig := globalping.Config{
		APIURL:    "https://api.globalping.io/v1",
		AuthURL:   "https://auth.globalping.io",
		UserAgent: fmt.Sprintf("doggo/%s (https://github.com/mr-karan/doggo)", buildVersion),
	}
	gpToken := os.Getenv("GLOBALPING_TOKEN")
	if gpToken != "" {
		gpConfig.AuthToken = &globalping.Token{
			AccessToken: gpToken,
			Expiry:      time.Now().Add(math.MaxInt64),
		}
	}
	globlpingClient := globalping.NewClient(gpConfig)

	app := app.New(logger, globlpingClient, buildVersion)

	if err := k.Unmarshal("", &app.QueryFlags); err != nil {
		logger.Error("Error loading args", "error", err)
		os.Exit(2)
	}

	loadNameservers(&app, cfg.flagSet.Args())
	return &app
}

func loadNameservers(app *app.App, args []string) {
	flagNameservers := k.Strings("nameserver")
	unparsedNameservers, qt, qc, qn := loadUnparsedArgs(args)

	if len(flagNameservers) > 0 {
		app.QueryFlags.Nameservers = flagNameservers
	} else {
		app.QueryFlags.Nameservers = unparsedNameservers
	}

	app.QueryFlags.QTypes = append(app.QueryFlags.QTypes, qt...)
	app.QueryFlags.QClasses = append(app.QueryFlags.QClasses, qc...)
	app.QueryFlags.QNames = append(app.QueryFlags.QNames, qn...)
}

func loadResolvers(app *app.App, cfg *config) ([]resolvers.Resolver, error) {
	return resolvers.LoadResolvers(resolvers.Options{
		Nameservers:        app.Nameservers,
		UseIPv4:            app.QueryFlags.UseIPv4,
		UseIPv6:            app.QueryFlags.UseIPv6,
		SearchList:         app.ResolverOpts.SearchList,
		Ndots:              app.ResolverOpts.Ndots,
		Timeout:            cfg.timeout,
		Logger:             app.Logger,
		Strategy:           app.QueryFlags.Strategy,
		InsecureSkipVerify: app.QueryFlags.InsecureSkipVerify,
		TLSHostname:        app.QueryFlags.TLSHostname,
	})
}

func performLookup(app *app.App, cfg *config) ([]resolvers.Response, []error) {
	ctx, cancel := context.WithTimeout(context.Background(), cfg.timeout)
	defer cancel()

	var (
		wg           sync.WaitGroup
		mu           sync.Mutex
		allResponses []resolvers.Response
		allErrors    []error
	)

	for _, resolver := range app.Resolvers {
		wg.Add(1)
		go func(r resolvers.Resolver) {
			defer wg.Done()
			responses, err := r.Lookup(ctx, app.Questions, cfg.queryFlags)
			mu.Lock()
			defer mu.Unlock()
			// Collect any responses the resolver produced even when err != nil
			// so partial successes within a single resolver still surface.
			allResponses = append(allResponses, responses...)
			if err != nil {
				allErrors = append(allErrors, &resolvers.LookupError{
					Nameserver: r.Address(),
					Err:        err,
				})
			}
		}(resolver)
	}

	wg.Wait()
	return allResponses, allErrors
}

func outputResults(app *app.App, responses []resolvers.Response, responseErrors []error) {
	if app.QueryFlags.ShowJSON {
		outputJSON(app.Logger, responses, responseErrors)
	} else {
		// Full failure: no resolver produced a usable response. Surface every
		// per-resolver error so the user can see which nameservers failed and
		// why, then exit with the legacy lookup-failure code.
		if len(responses) == 0 && len(responseErrors) > 0 {
			for _, err := range responseErrors {
				logResolverError(app.Logger, slog.LevelError, "Error looking up DNS records", err)
			}
			os.Exit(exitLookupFailure)
		}
		// Partial success: at least one resolver answered while another
		// failed. Demote the failure to a warning and print whatever we have.
		for _, err := range responseErrors {
			logResolverError(app.Logger, slog.LevelWarn, "lookup failed", err)
		}
		app.Output(responses)
	}

	if len(responseErrors) > 0 && len(responses) > 0 {
		os.Exit(exitPartialFailure)
	}
	if len(responseErrors) > 0 {
		os.Exit(exitLookupFailure)
	}
}

// logResolverError emits a per-resolver lookup error at the given level,
// unwrapping LookupError so the nameserver shows up as its own structured
// field rather than embedded in the message.
func logResolverError(logger *slog.Logger, level slog.Level, msg string, err error) {
	var lookupErr *resolvers.LookupError
	if errors.As(err, &lookupErr) {
		logger.Log(context.Background(), level, msg,
			"nameserver", lookupErr.Nameserver,
			"error", lookupErr.Err,
		)
		return
	}
	logger.Log(context.Background(), level, msg, "error", err)
}

// resolverErrorJSON is the per-resolver error shape returned in JSON output.
type resolverErrorJSON struct {
	Nameserver string `json:"nameserver,omitempty"`
	Error      string `json:"error"`
}

func outputJSON(logger *slog.Logger, responses []resolvers.Response, responseErrors []error) {
	jsonOutput := struct {
		Responses []resolvers.Response `json:"responses,omitempty"`
		Errors    []resolverErrorJSON  `json:"errors,omitempty"`
		// Error is kept for backwards compatibility with scripts that parsed
		// the previous schema. It is populated only on full failure.
		Error string `json:"error,omitempty"`
	}{
		Responses: responses,
	}

	for _, err := range responseErrors {
		var lookupErr *resolvers.LookupError
		if errors.As(err, &lookupErr) {
			jsonOutput.Errors = append(jsonOutput.Errors, resolverErrorJSON{
				Nameserver: lookupErr.Nameserver,
				Error:      lookupErr.Err.Error(),
			})
			continue
		}
		jsonOutput.Errors = append(jsonOutput.Errors, resolverErrorJSON{Error: err.Error()})
	}

	if len(responses) == 0 && len(responseErrors) > 0 {
		jsonOutput.Error = responseErrors[0].Error()
	}

	jsonData, err := json.MarshalIndent(jsonOutput, "", "  ")
	if err != nil {
		logger.Error("Error marshaling JSON")
		os.Exit(exitGenericFailure)
	}
	fmt.Println(string(jsonData))
}
