// Package whois provides WHOIS lookup functionality for ASN information.
package whois

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"net"
	"regexp"
	"strings"
	"time"
)

// Timeout constants for WHOIS queries.
const (
	dialTimeout  = 10 * time.Second
	readTimeout  = 30 * time.Second
	writeTimeout = 5 * time.Second
)

// Parsing constants.
const (
	keyValueParts       = 2 // Expected parts when splitting "key: value"
	lacnicDateFormatLen = 8 // Length of YYYYMMDD date format
)

// WHOIS server addresses.
const (
	whoisServerIANA    = "whois.iana.org:43"
	whoisServerARIN    = "whois.arin.net:43"
	whoisServerRIPE    = "whois.ripe.net:43"
	whoisServerAPNIC   = "whois.apnic.net:43"
	whoisServerLACNIC  = "whois.lacnic.net:43"
	whoisServerAFRINIC = "whois.afrinic.net:43"
)

// RIR identifiers.
const (
	RIRARIN   = "ARIN"
	RIRRIPE   = "RIPE"
	RIRAPNIC  = "APNIC"
	RIRLACNIC = "LACNIC"
	RIRAFRNIC = "AFRINIC"
)

// Errors LookupASN returns when a WHOIS answer has nothing to store.
var (
	// ErrNoReferral is returned when IANA refers the AS to no RIR, as it
	// does for a reserved, private-use or unallocated AS number.
	ErrNoReferral = errors.New("IANA refers the AS to no RIR")

	errNoASBlock = errors.New("IANA answered with neither a referral " +
		"nor an AS block")
	errNoFields = errors.New("RIR answered with no WHOIS fields")
)

// ASNInfo contains parsed WHOIS information for an ASN.
type ASNInfo struct {
	ASN         int
	ASName      string
	OrgName     string
	OrgID       string
	Address     string
	CountryCode string
	AbuseEmail  string
	AbusePhone  string
	TechEmail   string
	TechPhone   string
	RIR         string
	RegDate     *time.Time
	LastMod     *time.Time
	RawResponse string
}

// Client performs WHOIS lookups for ASNs.
type Client struct {
	// DialContext opens the connection to a WHOIS server, given as
	// host:port. NewClient sets a net.Dialer's; tests replace it to reach
	// local servers instead.
	DialContext func(ctx context.Context, network, address string) (net.Conn, error)
}

// NewClient creates a new WHOIS client.
func NewClient() *Client {
	dialer := &net.Dialer{Timeout: dialTimeout}

	return &Client{DialContext: dialer.DialContext}
}

// LookupASN asks IANA which RIR holds the given ASN, asks that RIR, and
// returns what the RIR's answer says. It returns ErrNoReferral when IANA
// refers the AS to no RIR, and another error when a query fails, when IANA's
// answer has neither a referral nor the AS's block, or when the RIR's answer
// carries none of ASNInfo's fields, as a refusal of service does.
func (c *Client) LookupASN(ctx context.Context, asn int) (*ASNInfo, error) {
	// Query IANA first to find the authoritative RIR
	query := fmt.Sprintf("AS%d", asn)

	ianaResp, err := c.query(ctx, whoisServerIANA, query)
	if err != nil {
		return nil, fmt.Errorf("IANA query failed: %w", err)
	}

	// Determine RIR from IANA response
	rir, whoisServer := c.parseIANAReferral(ianaResp)
	if whoisServer == "" {
		// IANA names the block of an AS it refers to no RIR. An answer
		// without that, such as an empty one, says nothing about the AS.
		if !strings.Contains(ianaResp, "as-block:") {
			return nil, fmt.Errorf("%s: %w", query, errNoASBlock)
		}

		return nil, fmt.Errorf("%s: %w", query, ErrNoReferral)
	}

	// Query the authoritative RIR
	rirResp, err := c.query(ctx, whoisServer, query)
	if err != nil {
		return nil, fmt.Errorf("RIR query failed: %w", err)
	}

	// The fields come from the RIR's answer alone: IANA's describes the
	// block the AS is in.
	info := c.parseResponse(asn, rir, rirResp)
	if info.isEmpty() {
		return nil, fmt.Errorf("%s from %s: %w", query, whoisServer, errNoFields)
	}

	info.RawResponse = ianaResp + "\n" + rirResp

	return info, nil
}

// isEmpty reports whether info holds none of the fields a WHOIS answer
// fills, as when the answer is only comment and error lines.
func (info *ASNInfo) isEmpty() bool {
	return info.ASName == "" && info.OrgName == "" && info.OrgID == "" &&
		info.Address == "" && info.CountryCode == "" &&
		info.AbuseEmail == "" && info.AbusePhone == "" &&
		info.TechEmail == "" && info.TechPhone == "" &&
		info.RegDate == nil && info.LastMod == nil
}

// query performs a raw WHOIS query to the specified server.
func (c *Client) query(ctx context.Context, server, query string) (string, error) {
	conn, err := c.DialContext(ctx, "tcp", server)
	if err != nil {
		return "", fmt.Errorf("dial %s: %w", server, err)
	}
	defer func() { _ = conn.Close() }()

	// Set deadlines
	err = conn.SetWriteDeadline(time.Now().Add(writeTimeout))
	if err != nil {
		return "", fmt.Errorf("set write deadline: %w", err)
	}

	// Send query
	_, err = fmt.Fprintf(conn, "%s\r\n", query)
	if err != nil {
		return "", fmt.Errorf("write query: %w", err)
	}

	// Read response
	err = conn.SetReadDeadline(time.Now().Add(readTimeout))
	if err != nil {
		return "", fmt.Errorf("set read deadline: %w", err)
	}

	var sb strings.Builder

	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		sb.WriteString(scanner.Text())
		sb.WriteString("\n")
	}

	err = scanner.Err()
	if err != nil {
		return sb.String(), fmt.Errorf("read response: %w", err)
	}

	return sb.String(), nil
}

// parseIANAReferral extracts the RIR and WHOIS server from an IANA response.
// The server is "" when the response has no referral.
func (c *Client) parseIANAReferral(response string) (string, string) {
	for line := range strings.SplitSeq(response, "\n") {
		line = strings.TrimSpace(line)

		// Look for "refer:" line
		if strings.HasPrefix(strings.ToLower(line), "refer:") {
			server := strings.TrimSpace(strings.TrimPrefix(line, "refer:"))
			server = strings.TrimSpace(strings.TrimPrefix(server, "Refer:"))

			switch {
			case strings.Contains(server, "arin"):
				return RIRARIN, whoisServerARIN
			case strings.Contains(server, "ripe"):
				return RIRRIPE, whoisServerRIPE
			case strings.Contains(server, "apnic"):
				return RIRAPNIC, whoisServerAPNIC
			case strings.Contains(server, "lacnic"):
				return RIRLACNIC, whoisServerLACNIC
			case strings.Contains(server, "afrinic"):
				return RIRAFRNIC, whoisServerAFRINIC
			default:
				// Unknown server, add port if missing
				if !strings.Contains(server, ":") {
					server += ":43"
				}

				return "", server
			}
		}
	}

	return "", ""
}

// parseResponse extracts ASN information from a WHOIS response.
func (c *Client) parseResponse(asn int, rir, response string) *ASNInfo {
	info := &ASNInfo{
		ASN:         asn,
		RIR:         rir,
		RawResponse: response,
	}

	lines := strings.Split(response, "\n")

	var addressLines []string

	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "%") ||
			strings.HasPrefix(line, "#") {
			continue
		}

		// Split on first colon
		parts := strings.SplitN(line, ":", keyValueParts)
		if len(parts) != keyValueParts {
			continue
		}

		key := strings.TrimSpace(strings.ToLower(parts[0]))
		value := strings.TrimSpace(parts[1])

		if value == "" {
			continue
		}

		// Address (collect multiple lines)
		if key == "address" {
			addressLines = append(addressLines, value)

			continue
		}

		c.setResponseField(info, key, value)
	}

	// Combine address lines
	if len(addressLines) > 0 {
		info.Address = strings.Join(addressLines, "\n")
	}

	// Extract abuse email from comment lines (common in ARIN responses)
	if info.AbuseEmail == "" {
		info.AbuseEmail = c.extractAbuseEmail(response)
	}

	return info
}

// setResponseField stores the value of one "key: value" line of a WHOIS
// response in info, keeping the first value seen for each field.
func (c *Client) setResponseField(info *ASNInfo, key, value string) {
	switch key {
	// Country
	case "country":
		if info.CountryCode == "" && len(value) == 2 {
			info.CountryCode = strings.ToUpper(value)
		}

	// Registration dates
	case "regdate", "created":
		if info.RegDate == nil {
			info.RegDate = c.parseDate(value)
		}
	case "updated", "last-modified", "changed":
		if info.LastMod == nil {
			info.LastMod = c.parseDate(value)
		}

	default:
		field := textFieldForKey(info, key)
		if field != nil && *field == "" {
			*field = value
		}
	}
}

// textFieldForKey returns the text field of info that a WHOIS response key
// fills, or nil when the key fills none.
func textFieldForKey(info *ASNInfo, key string) *string {
	switch key {
	// AS Name (varies by RIR)
	case "asname", "as-name":
		return &info.ASName

	// Organization
	case "orgname", "org-name", "owner":
		return &info.OrgName
	case "orgid", "org-id", "org":
		return &info.OrgID

	// Abuse contact
	case "orgabuseemail", "abuse-mailbox":
		return &info.AbuseEmail
	case "orgabusephone":
		return &info.AbusePhone

	// Tech contact
	case "orgtechemail":
		return &info.TechEmail
	case "orgtechphone":
		return &info.TechPhone

	default:
		return nil
	}
}

// parseDate attempts to parse various date formats used in WHOIS responses.
func (c *Client) parseDate(value string) *time.Time {
	// Common formats
	formats := []string{
		"2006-01-02",
		"2006-01-02T15:04:05Z",
		"2006-01-02T15:04:05-07:00",
		"20060102",
		"02-Jan-2006",
	}

	// Clean up value
	value = strings.TrimSpace(value)
	// Handle "YYYYMMDD" format from LACNIC
	if len(value) == lacnicDateFormatLen {
		t, err := time.Parse("20060102", value)
		if err == nil {
			return &t
		}
	}

	for _, format := range formats {
		t, err := time.Parse(format, value)
		if err == nil {
			return &t
		}
	}

	return nil
}

// extractAbuseEmail extracts abuse email from response using regex.
func (c *Client) extractAbuseEmail(response string) string {
	// Look for "Abuse contact for 'AS...' is 'email@domain'"
	re := regexp.MustCompile(`[Aa]buse contact.*?is\s+['"]?([^\s'"]+@[^\s'"]+)['"]?`)
	if matches := re.FindStringSubmatch(response); len(matches) > 1 {
		return matches[1]
	}

	return ""
}