//go:build generate

package main

import (
	"context"
	"encoding/json"
	"log/slog"
	"net/http"
	"os"
	"text/template"
	"time"

	"github.com/AdguardTeam/AdGuardDNS/internal/agdhttp"
	"github.com/AdguardTeam/AdGuardDNS/internal/geoip"
	"github.com/AdguardTeam/golibs/errors"
	"github.com/AdguardTeam/golibs/httphdr"
	"github.com/AdguardTeam/golibs/logutil/slogutil"
	"github.com/AdguardTeam/golibs/osutil"
)

func main() {
	ctx := context.Background()
	logger := slogutil.New(nil)
	defer slogutil.RecoverAndExit(ctx, logger, osutil.ExitCodeFailure)

	c := &http.Client{
		Timeout: 10 * time.Second,
	}

	req := errors.Must(http.NewRequestWithContext(ctx, http.MethodGet, countriesASNURL, nil))

	req.Header.Add(httphdr.UserAgent, agdhttp.UserAgent())

	resp := errors.Must(c.Do(req))
	defer slogutil.CloseAndLog(ctx, logger, resp.Body, slog.LevelError)

	errors.Check(agdhttp.CheckStatus(resp, http.StatusOK))

	out := errors.Must(os.OpenFile("./asntops.go", os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o664))
	defer slogutil.CloseAndLog(ctx, logger, out, slog.LevelError)

	defaultCountryTopASNs := map[geoip.Country][]geoip.ASN{}
	errors.Check(json.NewDecoder(resp.Body).Decode(&defaultCountryTopASNs))

	// Don't use a *container.MapSet here, because the map is iterated over in
	// the template.
	defaultTopASNs := map[geoip.ASN]struct{}{}
	for _, asns := range defaultCountryTopASNs {
		for _, asn := range asns {
			if asn != 0 {
				defaultTopASNs[asn] = struct{}{}
			}
		}
	}

	type templateData struct {
		DefaultTopASNs        map[geoip.ASN]struct{}
		DefaultCountryTopASNs map[geoip.Country][]geoip.ASN
	}

	tmplData := &templateData{
		DefaultTopASNs:        defaultTopASNs,
		DefaultCountryTopASNs: defaultCountryTopASNs,
	}

	tmpl := template.Must(template.New("main").Parse(tmplStr))
	errors.Check(tmpl.Execute(out, tmplData))
}

// countriesASNURL is the default URL to get the per-country top ASN statistics
// from.
const countriesASNURL = `https://static.adtidy.org/dns/countries_asn.json`

// tmplStr is the template of the generated Go code.
const tmplStr = `// Code generated by go run ./asntops_generate.go; DO NOT EDIT.

package geoip

import "github.com/AdguardTeam/golibs/container"

// DefaultTopASNs contains all specially handled ASNs.
var DefaultTopASNs = container.NewMapSet[ASN](
{{- range $asn, $_ := .DefaultTopASNs }}
	{{ $asn }},
{{- end }}
)

// DefaultCountryTopASNs is a mapping of a country to their top ASNs.
var DefaultCountryTopASNs = map[Country]ASN{
{{- range $ctry, $ASNs := .DefaultCountryTopASNs }}
{{- if gt (len $ASNs) 0 }}
	Country{{ $ctry }}: {{ index $ASNs 0 }},
{{- else }}
{{- continue }}
{{- end }}
{{- end }}
}
`
