package federation

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/dariuszjaniak/warp/server/internal/buildinfo"
	"github.com/dariuszjaniak/warp/server/internal/matrix/signing"
	"github.com/dariuszjaniak/warp/server/internal/tracing"
)

// Errors an outbound request can fail with before anything is sent.
var (
	// ErrDisabled means federation is switched off. Every outbound request
	// checks the switch itself, so that a component that missed the
	// toggle still cannot reach another server.
	ErrDisabled = errors.New("federation is disabled")

	// ErrNotAllowed means the destination is excluded by the allow and
	// deny lists.
	ErrNotAllowed = errors.New("federation with that server is not allowed")

	// ErrSelf means a request was addressed to this server. Asking
	// ourselves over the network can only be a bug in the caller, and on
	// a server whose name does not resolve to itself it would ask
	// somebody else.
	ErrSelf = errors.New("refusing to send a federation request to this server")
)

// HTTPError is a non-2xx answer from another server.
type HTTPError struct {
	Status  int
	Code    string
	Message string
}

func (e *HTTPError) Error() string {
	if e.Code != "" {
		return fmt.Sprintf("remote server answered %d %s: %s", e.Status, e.Code, e.Message)
	}
	return fmt.Sprintf("remote server answered %d", e.Status)
}

// IsNotFound reports a 404 from another server.
func IsNotFound(err error) bool {
	var httpErr *HTTPError
	return errors.As(err, &httpErr) && httpErr.Status == http.StatusNotFound
}

// Response size limits. A send_join answer for a large room is the room's
// entire state and auth chain, which for a public room of tens of
// thousands of members is hundreds of megabytes; nothing else comes close.
const (
	defaultMaxResponse = 16 << 20
	largeMaxResponse   = 1 << 30
)

// Request timeouts. Joining a large room makes the resident server
// assemble and sign that whole answer first, which takes minutes.
const (
	defaultRequestTimeout = time.Minute
	longRequestTimeout    = 10 * time.Minute
)

// ResolveFunc turns a server name into where to connect.
type ResolveFunc func(ctx context.Context, serverName string) (Resolved, error)

// ClientOptions configure a Client.
type ClientOptions struct {
	// ServerName is this server's name: the origin of every request.
	ServerName string

	// Key signs every request.
	Key *signing.KeyPair

	// Resolve is server discovery. Nil uses a Resolver.
	Resolve ResolveFunc

	// TLSConfig is the client TLS configuration, for tests and for
	// deployments with a private CA. Nil uses the system roots.
	TLSConfig *tls.Config

	// Enabled reports the live federation switch. Nil means off.
	Enabled func() bool

	// Policy filters destinations. Nil allows any.
	Policy *Policy

	// AllowPrivateMediaRedirects lets a media download follow another
	// server's redirect to a loopback, private or link-local address.
	// Only tests, whose servers are all on loopback, set it: anywhere
	// else it would let another server point this one at its own
	// internal network and read the answer back through a client.
	AllowPrivateMediaRedirects bool
}

// Client sends signed Server-Server API requests.
type Client struct {
	serverName string
	key        *signing.KeyPair
	resolve    ResolveFunc
	http       *http.Client
	enabled    func() bool
	policy     *Policy
	userAgent  string

	// mediaHTTP follows the redirects a media download may answer with.
	// It is not the federation transport: the URL is an ordinary web
	// address, often a CDN, and nothing is signed.
	mediaHTTP *http.Client
}

// dialTargetKey carries the resolved addresses of a request down to the
// dialer, which is the only place they are needed: the URL holds the Host
// the certificate is checked against, and the address the connection goes
// to can be something else entirely.
type dialTargetKey struct{}

// NewClient builds a client.
func NewClient(opts ClientOptions) *Client {
	dialer := &net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}
	tlsConfig := opts.TLSConfig
	if tlsConfig == nil {
		tlsConfig = &tls.Config{MinVersion: tls.VersionTLS12}
	}
	transport := &http.Transport{
		Proxy: nil,
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			targets, _ := ctx.Value(dialTargetKey{}).([]string)
			if len(targets) == 0 {
				return dialer.DialContext(ctx, network, addr)
			}
			var lastErr error
			for _, target := range targets {
				conn, err := dialer.DialContext(ctx, network, target)
				if err == nil {
					return conn, nil
				}
				lastErr = err
			}
			return nil, lastErr
		},
		TLSClientConfig:       tlsConfig,
		TLSHandshakeTimeout:   15 * time.Second,
		ResponseHeaderTimeout: 0, // bounded by the request context instead
		MaxIdleConns:          200,
		MaxIdleConnsPerHost:   4,
		IdleConnTimeout:       90 * time.Second,
		ForceAttemptHTTP2:     true,
	}

	if opts.Enabled == nil {
		opts.Enabled = func() bool { return false }
	}
	if opts.Policy == nil {
		opts.Policy = NewPolicy(nil, nil)
	}
	c := &Client{
		serverName: opts.ServerName,
		key:        opts.Key,
		resolve:    opts.Resolve,
		http: &http.Client{
			// Each request is a client span, and carries the trace on.
			Transport: tracing.Transport(transport),
			// Federation endpoints do not redirect; following one would
			// send a signed request somewhere it was not addressed.
			CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
		},
		enabled:   opts.Enabled,
		policy:    opts.Policy,
		userAgent: buildinfo.UserAgent(),
		mediaHTTP: newMediaRedirectClient(tlsConfig, opts.AllowPrivateMediaRedirects),
	}
	if c.resolve == nil {
		// The well-known document is fetched with ordinary HTTPS and the
		// same roots as everything else.
		resolver := NewResolver(&http.Client{
			Timeout:   30 * time.Second,
			Transport: &http.Transport{TLSClientConfig: tlsConfig, Proxy: nil},
		})
		c.resolve = resolver.Resolve
	}
	return c
}

// ServerName returns the origin this client signs as.
func (c *Client) ServerName() string { return c.serverName }

// check refuses a destination before anything goes on the wire.
func (c *Client) check(destination string) error {
	switch {
	case !c.enabled():
		return ErrDisabled
	case destination == c.serverName:
		return ErrSelf
	case !c.policy.Allowed(destination):
		return fmt.Errorf("%w: %s", ErrNotAllowed, destination)
	}
	return nil
}

// request describes one call.
type request struct {
	method      string
	destination string
	path        string // already escaped
	query       url.Values
	body        any
	signed      bool
	maxBytes    int64
	timeout     time.Duration
}

// do performs a request and decodes a JSON answer into out (which may be
// nil).
func (c *Client) do(ctx context.Context, req request, out any) error {
	resp, cancel, err := c.send(ctx, req)
	if err != nil {
		return err
	}
	defer cancel()
	defer func() { _ = resp.Body.Close() }()

	limit := req.maxBytes
	if limit == 0 {
		limit = defaultMaxResponse
	}
	body, err := io.ReadAll(io.LimitReader(resp.Body, limit+1))
	if err != nil {
		return fmt.Errorf("%s %s%s: read response: %w", req.method, req.destination, req.path, err)
	}
	if int64(len(body)) > limit {
		return fmt.Errorf("%s %s%s: response is larger than %d bytes", req.method, req.destination, req.path, limit)
	}

	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		return httpError(resp.StatusCode, body)
	}
	if out == nil {
		return nil
	}
	if err := json.Unmarshal(body, out); err != nil {
		return fmt.Errorf("%s %s%s: decode response: %w", req.method, req.destination, req.path, err)
	}
	return nil
}

// send performs a request and hands back the response with its body
// unread, for the callers that stream it. The caller closes the body and
// then calls cancel, which ends the request's deadline.
func (c *Client) send(ctx context.Context, req request) (*http.Response, context.CancelFunc, error) {
	if err := c.check(req.destination); err != nil {
		return nil, nil, err
	}
	resolved, err := c.resolve(ctx, req.destination)
	if err != nil {
		return nil, nil, fmt.Errorf("resolve %s: %w", req.destination, err)
	}

	timeout := req.timeout
	if timeout == 0 {
		timeout = defaultRequestTimeout
	}
	ctx, cancel := context.WithTimeout(ctx, timeout)
	ctx = context.WithValue(ctx, dialTargetKey{}, resolved.Addresses)

	var content []byte
	if req.body != nil {
		if content, err = json.Marshal(req.body); err != nil {
			cancel()
			return nil, nil, fmt.Errorf("encode request: %w", err)
		}
	}

	target := "https://" + resolved.Host + req.path
	if len(req.query) > 0 {
		target += "?" + req.query.Encode()
	}
	httpReq, err := http.NewRequestWithContext(ctx, req.method, target, bytes.NewReader(content))
	if err != nil {
		cancel()
		return nil, nil, fmt.Errorf("build request: %w", err)
	}
	if content == nil {
		httpReq.Body = http.NoBody
	} else {
		httpReq.Header.Set("Content-Type", "application/json")
	}
	httpReq.Header.Set("User-Agent", c.userAgent)
	httpReq.Header.Set("Accept", "application/json")

	if req.signed {
		// The signature covers the request target exactly as it goes on
		// the wire, so it is taken from the built request rather than
		// from the path this function was handed.
		x, err := SignRequest(c.key, c.serverName, req.destination, req.method,
			httpReq.URL.RequestURI(), content)
		if err != nil {
			cancel()
			return nil, nil, fmt.Errorf("sign request: %w", err)
		}
		httpReq.Header.Set("Authorization", x.String())
	}

	resp, err := c.http.Do(httpReq)
	if err != nil {
		cancel()
		return nil, nil, fmt.Errorf("%s %s%s: %w", req.method, req.destination, req.path, err)
	}
	return resp, cancel, nil
}

// httpError turns a non-2xx answer into an *HTTPError, with the Matrix
// error envelope if the body was one.
func httpError(status int, body []byte) *HTTPError {
	httpErr := &HTTPError{Status: status}
	var envelope struct {
		Code    string `json:"errcode"`
		Message string `json:"error"`
	}
	if json.Unmarshal(body, &envelope) == nil {
		httpErr.Code, httpErr.Message = envelope.Code, envelope.Message
	}
	return httpErr
}

// esc escapes one path segment. Room, event and user IDs carry sigils and
// colons, which are legal in a path; slashes would not be, and an event ID
// of room versions 3 and before can contain one.
func esc(s string) string { return url.PathEscape(s) }

// ServerKeys fetches a server's own key response, unsigned and unverified:
// verifying it is the key ring's job.
func (c *Client) ServerKeys(ctx context.Context, serverName string) (json.RawMessage, error) {
	var out json.RawMessage
	err := c.do(ctx, request{method: http.MethodGet, destination: serverName, path: "/_matrix/key/v2/server"}, &out)
	return out, err
}

// QueryKeys asks a notary for other servers' keys.
func (c *Client) QueryKeys(ctx context.Context, notary string, query map[string]map[string]KeyCriteria) ([]json.RawMessage, error) {
	var out struct {
		ServerKeys []json.RawMessage `json:"server_keys"`
	}
	err := c.do(ctx, request{
		method: http.MethodPost, destination: notary, path: "/_matrix/key/v2/query",
		body: map[string]any{"server_keys": query},
	}, &out)
	return out.ServerKeys, err
}

// KeyCriteria is one key asked of a notary.
type KeyCriteria struct {
	MinimumValidUntilTS int64 `json:"minimum_valid_until_ts,omitempty"`
}

// SendTransaction delivers a transaction.
func (c *Client) SendTransaction(ctx context.Context, destination, txnID string, txn Transaction) (SendTransactionResponse, error) {
	var out SendTransactionResponse
	err := c.do(ctx, request{
		method: http.MethodPut, destination: destination, signed: true,
		path: "/_matrix/federation/v1/send/" + esc(txnID), body: txn,
	}, &out)
	return out, err
}

// MakeJoin asks a resident server for a join event template.
func (c *Client) MakeJoin(ctx context.Context, destination, roomID, userID string, versions []string) (MakeJoinResponse, error) {
	query := url.Values{}
	for _, v := range versions {
		query.Add("ver", v)
	}
	var out MakeJoinResponse
	err := c.do(ctx, request{
		method: http.MethodGet, destination: destination, signed: true,
		path:  "/_matrix/federation/v1/make_join/" + esc(roomID) + "/" + esc(userID),
		query: query,
	}, &out)
	return out, err
}

// SendJoin submits a signed join and receives the room's state.
func (c *Client) SendJoin(ctx context.Context, destination, roomID, eventID string, joinEvent json.RawMessage) (SendJoinResponse, error) {
	var out SendJoinResponse
	err := c.do(ctx, request{
		method: http.MethodPut, destination: destination, signed: true,
		path:     "/_matrix/federation/v2/send_join/" + esc(roomID) + "/" + esc(eventID),
		body:     joinEvent,
		maxBytes: largeMaxResponse,
		timeout:  longRequestTimeout,
	}, &out)
	return out, err
}

// MakeLeave asks a resident server for a leave event template.
func (c *Client) MakeLeave(ctx context.Context, destination, roomID, userID string) (MakeJoinResponse, error) {
	var out MakeJoinResponse
	err := c.do(ctx, request{
		method: http.MethodGet, destination: destination, signed: true,
		path: "/_matrix/federation/v1/make_leave/" + esc(roomID) + "/" + esc(userID),
	}, &out)
	return out, err
}

// SendLeave submits a signed leave.
func (c *Client) SendLeave(ctx context.Context, destination, roomID, eventID string, leaveEvent json.RawMessage) error {
	return c.do(ctx, request{
		method: http.MethodPut, destination: destination, signed: true,
		path: "/_matrix/federation/v2/send_leave/" + esc(roomID) + "/" + esc(eventID),
		body: leaveEvent,
	}, nil)
}

// Invite asks the invitee's server to countersign an invite.
func (c *Client) Invite(ctx context.Context, destination, roomID, eventID string, req InviteRequest) (InviteResponse, error) {
	var out InviteResponse
	err := c.do(ctx, request{
		method: http.MethodPut, destination: destination, signed: true,
		path: "/_matrix/federation/v2/invite/" + esc(roomID) + "/" + esc(eventID),
		body: req,
	}, &out)
	return out, err
}

// Event fetches one event.
func (c *Client) Event(ctx context.Context, destination, eventID string) (Transaction, error) {
	var out Transaction
	err := c.do(ctx, request{
		method: http.MethodGet, destination: destination, signed: true,
		path: "/_matrix/federation/v1/event/" + esc(eventID),
	}, &out)
	return out, err
}

// EventAuth fetches the auth chain of an event.
func (c *Client) EventAuth(ctx context.Context, destination, roomID, eventID string) (EventAuthResponse, error) {
	var out EventAuthResponse
	err := c.do(ctx, request{
		method: http.MethodGet, destination: destination, signed: true,
		path:     "/_matrix/federation/v1/event_auth/" + esc(roomID) + "/" + esc(eventID),
		maxBytes: largeMaxResponse,
		timeout:  longRequestTimeout,
	}, &out)
	return out, err
}

// StateIDs fetches the IDs of the room state at an event.
func (c *Client) StateIDs(ctx context.Context, destination, roomID, eventID string) (StateIDsResponse, error) {
	var out StateIDsResponse
	err := c.do(ctx, request{
		method: http.MethodGet, destination: destination, signed: true,
		path:     "/_matrix/federation/v1/state_ids/" + esc(roomID),
		query:    url.Values{"event_id": {eventID}},
		maxBytes: largeMaxResponse,
		timeout:  longRequestTimeout,
	}, &out)
	return out, err
}

// State fetches the room state at an event, events and all.
func (c *Client) State(ctx context.Context, destination, roomID, eventID string) (StateResponse, error) {
	var out StateResponse
	err := c.do(ctx, request{
		method: http.MethodGet, destination: destination, signed: true,
		path:     "/_matrix/federation/v1/state/" + esc(roomID),
		query:    url.Values{"event_id": {eventID}},
		maxBytes: largeMaxResponse,
		timeout:  longRequestTimeout,
	}, &out)
	return out, err
}

// MissingEvents asks for the events between what this server has and an
// event it was sent.
func (c *Client) MissingEvents(ctx context.Context, destination, roomID string, req MissingEventsRequest) (MissingEventsResponse, error) {
	var out MissingEventsResponse
	err := c.do(ctx, request{
		method: http.MethodPost, destination: destination, signed: true,
		path: "/_matrix/federation/v1/get_missing_events/" + esc(roomID),
		body: req,
	}, &out)
	return out, err
}

// Backfill asks for history before the given events.
func (c *Client) Backfill(ctx context.Context, destination, roomID string, from []string, limit int) (Transaction, error) {
	query := url.Values{"limit": {strconv.Itoa(limit)}}
	for _, id := range from {
		query.Add("v", id)
	}
	var out Transaction
	err := c.do(ctx, request{
		method: http.MethodGet, destination: destination, signed: true,
		path:  "/_matrix/federation/v1/backfill/" + esc(roomID),
		query: query,
	}, &out)
	return out, err
}

// QueryProfile asks a server for one of its users' profile. An empty field
// asks for the whole profile.
func (c *Client) QueryProfile(ctx context.Context, destination, userID, field string) (map[string]any, error) {
	query := url.Values{"user_id": {userID}}
	if field != "" {
		query.Set("field", field)
	}
	var out map[string]any
	err := c.do(ctx, request{
		method: http.MethodGet, destination: destination, signed: true,
		path: "/_matrix/federation/v1/query/profile", query: query,
	}, &out)
	return out, err
}

// QueryDirectory resolves an alias on the server that owns it.
func (c *Client) QueryDirectory(ctx context.Context, destination, alias string) (DirectoryResponse, error) {
	var out DirectoryResponse
	err := c.do(ctx, request{
		method: http.MethodGet, destination: destination, signed: true,
		path: "/_matrix/federation/v1/query/directory", query: url.Values{"room_alias": {alias}},
	}, &out)
	return out, err
}

// PublicRooms reads another server's public room directory.
func (c *Client) PublicRooms(ctx context.Context, destination string, limit int, since, search string) (json.RawMessage, error) {
	body := map[string]any{}
	if limit > 0 {
		body["limit"] = limit
	}
	if since != "" {
		body["since"] = since
	}
	if search = strings.TrimSpace(search); search != "" {
		body["filter"] = map[string]string{"generic_search_term": search}
	}
	var out json.RawMessage
	err := c.do(ctx, request{
		method: http.MethodPost, destination: destination, signed: true,
		path: "/_matrix/federation/v1/publicRooms", body: body,
	}, &out)
	return out, err
}

// Hierarchy asks a server about one room of a space and its children.
func (c *Client) Hierarchy(ctx context.Context, destination, roomID string, suggestedOnly bool) (json.RawMessage, error) {
	query := url.Values{}
	if suggestedOnly {
		query.Set("suggested_only", "true")
	}
	var out json.RawMessage
	err := c.do(ctx, request{
		method: http.MethodGet, destination: destination, signed: true,
		path: "/_matrix/federation/v1/hierarchy/" + url.PathEscape(roomID), query: query,
	}, &out)
	return out, err
}

// Version asks a server what software it runs.
func (c *Client) Version(ctx context.Context, destination string) (VersionResponse, error) {
	var out VersionResponse
	err := c.do(ctx, request{
		method: http.MethodGet, destination: destination,
		path: "/_matrix/federation/v1/version",
	}, &out)
	return out, err
}