package federation

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Server discovery: turning a server name into somewhere to connect.
//
// A Matrix server name is not an address. "example.org" may be served from
// matrix.example.org:443 because example.org says so in a well-known
// document, or from whatever an SRV record names, or from example.org:8448
// when neither exists. The spec's algorithm decides, in a fixed order, and
// the result has two halves that must not be confused: the address to dial,
// and the name the TLS certificate must be valid for (which is also the
// Host header). An SRV target is only ever the first: delegating by SRV
// does not let the target's certificate stand in for the server's.

// DefaultFederationPort is the port a server name without one is served on
// when nothing delegates it elsewhere.
const DefaultFederationPort = 8448

// ErrInvalidServerName is a server name that does not parse.
var ErrInvalidServerName = errors.New("invalid server name")

// ServerName is a parsed server name: a host, which is a DNS name or an IP
// literal, and an optional port.
type ServerName struct {
	Host string
	Port int

	// IsIP reports that Host is an IP literal, without the brackets an
	// IPv6 literal is written with.
	IsIP bool
}

// ParseServerName parses "host", "host:port", "1.2.3.4:port" or
// "[::1]:port".
func ParseServerName(s string) (ServerName, error) {
	if s == "" || len(s) > 255 {
		return ServerName{}, fmt.Errorf("%w: %q", ErrInvalidServerName, s)
	}

	host, portStr := s, ""
	if strings.HasPrefix(s, "[") {
		end := strings.IndexByte(s, ']')
		if end < 0 {
			return ServerName{}, fmt.Errorf("%w: %q has an unterminated IPv6 literal", ErrInvalidServerName, s)
		}
		host = s[1:end]
		rest := s[end+1:]
		if rest != "" {
			if !strings.HasPrefix(rest, ":") {
				return ServerName{}, fmt.Errorf("%w: %q", ErrInvalidServerName, s)
			}
			portStr = rest[1:]
		}
		ip := net.ParseIP(host)
		if ip == nil || ip.To4() != nil {
			return ServerName{}, fmt.Errorf("%w: %q is not an IPv6 literal", ErrInvalidServerName, s)
		}
	} else if i := strings.LastIndexByte(s, ':'); i >= 0 {
		host, portStr = s[:i], s[i+1:]
	}

	name := ServerName{Host: host}
	if portStr != "" || strings.HasSuffix(s, ":") {
		port, err := strconv.Atoi(portStr)
		if err != nil || port < 1 || port > 65535 {
			return ServerName{}, fmt.Errorf("%w: %q has a bad port", ErrInvalidServerName, s)
		}
		name.Port = port
	}

	if ip := net.ParseIP(host); ip != nil {
		if ip.To4() == nil && !strings.HasPrefix(s, "[") {
			// A bare IPv6 address with a port is ambiguous, which is why
			// the grammar requires the brackets.
			return ServerName{}, fmt.Errorf("%w: %q: an IPv6 literal must be bracketed", ErrInvalidServerName, s)
		}
		name.IsIP = true
		return name, nil
	}
	if !validDNSName(host) {
		return ServerName{}, fmt.Errorf("%w: %q is not a DNS name", ErrInvalidServerName, s)
	}
	return name, nil
}

// validDNSName applies the spec's dns-name grammar: letters, digits, '-'
// and '.', at most 255 characters.
func validDNSName(s string) bool {
	if s == "" || len(s) > 255 {
		return false
	}
	for _, r := range s {
		switch {
		case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '-', r == '.':
		default:
			return false
		}
	}
	return true
}

// hostPort renders a host and port for dialling, bracketing IPv6.
func hostPort(host string, port int) string {
	return net.JoinHostPort(host, strconv.Itoa(port))
}

// Resolved is where to send a server's federation traffic.
type Resolved struct {
	// Addresses are host:port pairs to try in order.
	Addresses []string

	// Host is the Host header and, without its port, the name the TLS
	// certificate must be valid for.
	Host string
}

// TLSServerName is the name the server's certificate must match.
func (r Resolved) TLSServerName() string {
	if host, _, err := net.SplitHostPort(r.Host); err == nil {
		return host
	}
	return strings.Trim(r.Host, "[]")
}

// Resolver implements server discovery, caching well-known lookups.
type Resolver struct {
	// HTTPClient fetches /.well-known/matrix/server. It should verify TLS
	// normally: the well-known document is fetched from the server name
	// itself, over ordinary HTTPS.
	HTTPClient *http.Client

	// LookupSRV is net.DefaultResolver.LookupSRV unless a test replaces it.
	LookupSRV func(ctx context.Context, service, proto, name string) (string, []*net.SRV, error)

	Now func() time.Time

	mu    sync.Mutex
	cache map[string]wellKnownEntry
}

type wellKnownEntry struct {
	server  string // the delegated m.server, or "" when there is none
	expires time.Time
}

// Well-known cache lifetimes, from the spec's guidance: a successful
// lookup is cached for what Cache-Control says, within these bounds, and a
// failed one for an hour, so that a server with no well-known document is
// not asked for one on every request.
const (
	wellKnownDefaultTTL = 24 * time.Hour
	wellKnownMaxTTL     = 48 * time.Hour
	wellKnownErrorTTL   = time.Hour

	// wellKnownMaxBytes caps the document. A real one is a line of JSON.
	wellKnownMaxBytes = 64 << 10
)

// NewResolver builds a resolver around an HTTP client.
func NewResolver(client *http.Client) *Resolver {
	if client == nil {
		client = &http.Client{Timeout: 30 * time.Second}
	}
	return &Resolver{
		HTTPClient: client,
		LookupSRV:  net.DefaultResolver.LookupSRV,
		Now:        time.Now,
		cache:      map[string]wellKnownEntry{},
	}
}

// Resolve runs the discovery algorithm for a server name.
func (r *Resolver) Resolve(ctx context.Context, serverName string) (Resolved, error) {
	name, err := ParseServerName(serverName)
	if err != nil {
		return Resolved{}, err
	}

	// Step 1: an IP literal is its own address.
	if name.IsIP {
		port := name.Port
		if port == 0 {
			port = DefaultFederationPort
		}
		return Resolved{Addresses: []string{hostPort(name.Host, port)}, Host: serverName}, nil
	}
	// Step 2: an explicit port means no delegation of any kind.
	if name.Port != 0 {
		return Resolved{Addresses: []string{hostPort(name.Host, name.Port)}, Host: serverName}, nil
	}

	// Step 3: a well-known document can delegate to another name.
	if delegated := r.wellKnown(ctx, name.Host); delegated != "" {
		dname, err := ParseServerName(delegated)
		if err == nil {
			switch {
			case dname.IsIP:
				port := dname.Port
				if port == 0 {
					port = DefaultFederationPort
				}
				return Resolved{Addresses: []string{hostPort(dname.Host, port)}, Host: delegated}, nil
			case dname.Port != 0:
				return Resolved{Addresses: []string{hostPort(dname.Host, dname.Port)}, Host: delegated}, nil
			default:
				return r.srvOrDefault(ctx, dname.Host), nil
			}
		}
		// A malformed delegation is treated as no delegation, which is
		// what the spec asks for an invalid well-known response.
	}

	// Steps 4 to 6: SRV records, then the default port.
	return r.srvOrDefault(ctx, name.Host), nil
}

// srvOrDefault looks up the SRV records of a hostname, the current service
// name first and the deprecated one second, falling back to port 8448.
// The Host header stays the hostname whatever the SRV target is.
func (r *Resolver) srvOrDefault(ctx context.Context, host string) Resolved {
	for _, service := range []string{"matrix-fed", "matrix"} {
		_, records, err := r.LookupSRV(ctx, service, "tcp", host)
		if err != nil || len(records) == 0 {
			continue
		}
		var addrs []string
		for _, rec := range records {
			target := strings.TrimSuffix(rec.Target, ".")
			if target == "" {
				continue
			}
			addrs = append(addrs, hostPort(target, int(rec.Port)))
		}
		if len(addrs) > 0 {
			return Resolved{Addresses: addrs, Host: host}
		}
	}
	return Resolved{Addresses: []string{hostPort(host, DefaultFederationPort)}, Host: host}
}

// wellKnown returns the delegated server name for a host, or "".
func (r *Resolver) wellKnown(ctx context.Context, host string) string {
	now := r.Now()
	r.mu.Lock()
	if entry, ok := r.cache[host]; ok && now.Before(entry.expires) {
		r.mu.Unlock()
		return entry.server
	}
	r.mu.Unlock()

	server, ttl := r.fetchWellKnown(ctx, host)

	r.mu.Lock()
	r.cache[host] = wellKnownEntry{server: server, expires: now.Add(ttl)}
	r.mu.Unlock()
	return server
}

// fetchWellKnown fetches and parses the document, returning the
// delegation and how long to believe it.
func (r *Resolver) fetchWellKnown(ctx context.Context, host string) (string, time.Duration) {
	ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
	defer cancel()

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://"+host+"/.well-known/matrix/server", nil)
	if err != nil {
		return "", wellKnownErrorTTL
	}
	resp, err := r.HTTPClient.Do(req)
	if err != nil {
		return "", wellKnownErrorTTL
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode != http.StatusOK {
		return "", wellKnownErrorTTL
	}
	body, err := io.ReadAll(io.LimitReader(resp.Body, wellKnownMaxBytes))
	if err != nil {
		return "", wellKnownErrorTTL
	}
	var doc struct {
		Server string `json:"m.server"`
	}
	if err := json.Unmarshal(body, &doc); err != nil || doc.Server == "" {
		return "", wellKnownErrorTTL
	}
	return doc.Server, cacheTTL(resp.Header.Get("Cache-Control"))
}

// cacheTTL reads max-age from a Cache-Control header, within the bounds
// the spec sets for a well-known lookup.
func cacheTTL(header string) time.Duration {
	for _, directive := range strings.Split(header, ",") {
		name, value, ok := strings.Cut(strings.TrimSpace(directive), "=")
		if !ok || !strings.EqualFold(name, "max-age") {
			continue
		}
		seconds, err := strconv.Atoi(strings.Trim(value, `"`))
		if err != nil || seconds < 0 {
			break
		}
		ttl := time.Duration(seconds) * time.Second
		if ttl > wellKnownMaxTTL {
			ttl = wellKnownMaxTTL
		}
		return ttl
	}
	return wellKnownDefaultTTL
}