package federation

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"syscall"
	"time"
)

// MediaResponse is media fetched from another server. The caller reads
// Body and closes it.
type MediaResponse struct {
	ContentType        string
	ContentDisposition string
	Body               io.ReadCloser
}

// ThumbnailRequest is the size of thumbnail asked of another server.
type ThumbnailRequest struct {
	Width    int
	Height   int
	Method   string
	Animated bool
}

// mediaTimeout bounds a media download, body included: fifty megabytes
// over a slow link is minutes, not seconds.
const mediaTimeout = 10 * time.Minute

// maxMediaMetadata caps the JSON part of a multipart media answer, which
// the spec defines as an empty object for now.
const maxMediaMetadata = 64 << 10

// errCodeUnrecognized is the answer of a server that predates an endpoint.
const errCodeUnrecognized = "M_UNRECOGNIZED"

// DownloadMedia fetches media from the server it was uploaded to.
//
// It asks the authenticated endpoint of Matrix 1.11 first. A server older
// than that answers 404 M_UNRECOGNIZED, and is then asked through the
// deprecated Client-Server download endpoint with allow_remote=false,
// which is what every server before 1.11 served media to other servers
// through.
func (c *Client) DownloadMedia(ctx context.Context, origin, mediaID string, wait time.Duration) (MediaResponse, error) {
	query := url.Values{}
	if wait > 0 {
		query.Set("timeout_ms", strconv.FormatInt(wait.Milliseconds(), 10))
	}
	legacy := url.Values{"allow_remote": {"false"}}
	if wait > 0 {
		legacy.Set("timeout_ms", strconv.FormatInt(wait.Milliseconds(), 10))
	}
	return c.fetchMedia(ctx, origin,
		"/_matrix/federation/v1/media/download/"+esc(mediaID), query,
		"/_matrix/media/v3/download/"+esc(origin)+"/"+esc(mediaID), legacy)
}

// MediaThumbnail fetches a thumbnail from the server the media was
// uploaded to, with the same fallback as DownloadMedia.
func (c *Client) MediaThumbnail(ctx context.Context, origin, mediaID string, req ThumbnailRequest, wait time.Duration) (MediaResponse, error) {
	query := url.Values{
		"width":  {strconv.Itoa(req.Width)},
		"height": {strconv.Itoa(req.Height)},
		"method": {req.Method},
	}
	if req.Animated {
		query.Set("animated", "true")
	}
	if wait > 0 {
		query.Set("timeout_ms", strconv.FormatInt(wait.Milliseconds(), 10))
	}
	legacy := url.Values{}
	for k, v := range query {
		legacy[k] = v
	}
	legacy.Set("allow_remote", "false")
	return c.fetchMedia(ctx, origin,
		"/_matrix/federation/v1/media/thumbnail/"+esc(mediaID), query,
		"/_matrix/media/v3/thumbnail/"+esc(origin)+"/"+esc(mediaID), legacy)
}

// fetchMedia asks the federation endpoint, and the legacy one if the
// server does not know the federation one.
func (c *Client) fetchMedia(ctx context.Context, origin, path string, query url.Values, legacyPath string, legacyQuery url.Values) (MediaResponse, error) {
	resp, cancel, err := c.send(ctx, request{
		method: http.MethodGet, destination: origin, signed: true,
		path: path, query: query, timeout: mediaTimeout,
	})
	if err != nil {
		return MediaResponse{}, err
	}
	if resp.StatusCode == http.StatusOK {
		media, err := c.readMultipart(ctx, resp, cancel)
		if err != nil {
			return MediaResponse{}, fmt.Errorf("GET %s%s: %w", origin, path, err)
		}
		return media, nil
	}

	failure := readFailure(resp, cancel)
	if failure.Status != http.StatusNotFound || failure.Code != errCodeUnrecognized {
		return MediaResponse{}, failure
	}

	resp, cancel, err = c.send(ctx, request{
		method: http.MethodGet, destination: origin,
		path: legacyPath, query: legacyQuery, timeout: mediaTimeout,
	})
	if err != nil {
		return MediaResponse{}, err
	}
	if resp.StatusCode != http.StatusOK {
		return MediaResponse{}, readFailure(resp, cancel)
	}
	return MediaResponse{
		ContentType:        resp.Header.Get("Content-Type"),
		ContentDisposition: resp.Header.Get("Content-Disposition"),
		Body:               &closer{Reader: resp.Body, close: closeAll(resp.Body, cancel)},
	}, nil
}

// readFailure reads a non-2xx answer and releases the request.
func readFailure(resp *http.Response, cancel context.CancelFunc) *HTTPError {
	body, _ := io.ReadAll(io.LimitReader(resp.Body, maxMediaMetadata))
	_ = resp.Body.Close()
	cancel()
	return httpError(resp.StatusCode, body)
}

// readMultipart takes apart the multipart/mixed answer of the federation
// media endpoints: a JSON part, which is metadata and empty for now, then
// either the content itself or a part with a Location header saying where
// to fetch it from.
func (c *Client) readMultipart(ctx context.Context, resp *http.Response, cancel context.CancelFunc) (MediaResponse, error) {
	release := closeAll(resp.Body, cancel)
	mediaType, params, err := mime.ParseMediaType(resp.Header.Get("Content-Type"))
	if err != nil || mediaType != "multipart/mixed" || params["boundary"] == "" {
		_ = release()
		return MediaResponse{}, fmt.Errorf("answer is %q, not multipart/mixed", resp.Header.Get("Content-Type"))
	}
	parts := multipart.NewReader(resp.Body, params["boundary"])

	metadata, err := parts.NextPart()
	if err != nil {
		_ = release()
		return MediaResponse{}, fmt.Errorf("read the metadata part: %w", err)
	}
	if _, err := io.Copy(io.Discard, io.LimitReader(metadata, maxMediaMetadata)); err != nil {
		_ = release()
		return MediaResponse{}, fmt.Errorf("read the metadata part: %w", err)
	}

	content, err := parts.NextPart()
	if err != nil {
		_ = release()
		return MediaResponse{}, fmt.Errorf("read the content part: %w", err)
	}
	if location := content.Header.Get("Location"); location != "" {
		_ = release()
		return c.followMediaRedirect(ctx, location)
	}
	return MediaResponse{
		ContentType:        content.Header.Get("Content-Type"),
		ContentDisposition: content.Header.Get("Content-Disposition"),
		Body:               &closer{Reader: content, close: release},
	}, nil
}

// followMediaRedirect fetches content another server said is elsewhere.
func (c *Client) followMediaRedirect(ctx context.Context, location string) (MediaResponse, error) {
	target, err := url.Parse(location)
	if err != nil || (target.Scheme != "https" && target.Scheme != "http") || target.Host == "" {
		return MediaResponse{}, fmt.Errorf("media redirect to %q is not a web address", location)
	}
	ctx, cancel := context.WithTimeout(ctx, mediaTimeout)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, target.String(), http.NoBody)
	if err != nil {
		cancel()
		return MediaResponse{}, fmt.Errorf("build the media redirect request: %w", err)
	}
	req.Header.Set("User-Agent", c.userAgent)
	resp, err := c.mediaHTTP.Do(req)
	if err != nil {
		cancel()
		return MediaResponse{}, fmt.Errorf("follow the media redirect: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		failure := readFailure(resp, cancel)
		return MediaResponse{}, fmt.Errorf("follow the media redirect: %w", failure)
	}
	return MediaResponse{
		ContentType:        resp.Header.Get("Content-Type"),
		ContentDisposition: resp.Header.Get("Content-Disposition"),
		Body:               &closer{Reader: resp.Body, close: closeAll(resp.Body, cancel)},
	}, nil
}

// errPrivateAddress refuses a media redirect into a private network.
var errPrivateAddress = errors.New("refusing to fetch media from a private address")

// newMediaRedirectClient builds the client media redirects are followed
// with. Its dialer checks the address it actually connects to — after DNS,
// on every redirect — so a name that resolves into a private network is
// refused as surely as an IP literal.
func newMediaRedirectClient(tlsConfig *tls.Config, allowPrivate bool) *http.Client {
	dialer := &net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}
	if !allowPrivate {
		dialer.Control = func(_, address string, _ syscall.RawConn) error {
			host, _, err := net.SplitHostPort(address)
			if err != nil {
				return err
			}
			ip := net.ParseIP(host)
			if ip == nil || !isPublic(ip) {
				return fmt.Errorf("%w: %s", errPrivateAddress, host)
			}
			return nil
		}
	}
	return &http.Client{
		Transport: &http.Transport{
			Proxy:               nil,
			DialContext:         dialer.DialContext,
			TLSClientConfig:     tlsConfig,
			TLSHandshakeTimeout: 15 * time.Second,
			MaxIdleConns:        20,
			IdleConnTimeout:     90 * time.Second,
		},
		CheckRedirect: func(_ *http.Request, via []*http.Request) error {
			if len(via) >= 5 {
				return errors.New("too many media redirects")
			}
			return nil
		},
	}
}

// isPublic reports an address on the public internet.
func isPublic(ip net.IP) bool {
	return !ip.IsLoopback() && !ip.IsPrivate() && !ip.IsUnspecified() &&
		!ip.IsLinkLocalUnicast() && !ip.IsLinkLocalMulticast() &&
		!ip.IsInterfaceLocalMulticast() && !ip.IsMulticast()
}

// closer is a reader whose Close releases more than the reader.
type closer struct {
	io.Reader
	close func() error
}

func (c *closer) Close() error { return c.close() }

// closeAll closes a response body and ends its request's deadline.
func closeAll(body io.Closer, cancel context.CancelFunc) func() error {
	return func() error {
		err := body.Close()
		cancel()
		return err
	}
}