// Copyright 2018 The yeqown Author. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package proxy

import (
	"errors"
	"net"
	"net/http"

	"github.com/valyala/fasthttp"
)

const _fasthttpHostClientName = "reverse-proxy"

// ReverseProxy reverse handler using fasthttp.HostClient
type ReverseProxy struct {
	// bla keeps balancer instance
	bla IBalancer

	// clients
	clients []*fasthttp.HostClient

	// opt contains finally option to open reverseProxy
	opt *buildOption
}

// NewReverseProxyWith create an ReverseProxy with options
func NewReverseProxyWith(options ...Option) (*ReverseProxy, error) {
	option := defaultBuildOption()
	for _, opt := range options {
		opt.apply(option)
	}

	proxy := &ReverseProxy{
		bla:     nil,
		opt:     option,
		clients: make([]*fasthttp.HostClient, 0, 2),
	}

	err := proxy.init()
	if err != nil {
		return nil, err
	}

	return proxy, nil
}

// init initialize the ReverseProxy with options,
// if opted.OpenBalance is true then create a balancer to ReverseProxy
// else just create a HostClient to ReverseProxy and use it.
func (p *ReverseProxy) init() error {
	if len(p.opt.addresses) == 0 {
		return errors.New("no upstream server address")
	}

	if p.opt.openBalance {
		// config balancer
		p.clients = make([]*fasthttp.HostClient, 0, len(p.opt.addresses))
		p.bla = NewBalancer(p.opt.weights)

		for _, addr := range p.opt.addresses {
			client := &fasthttp.HostClient{
				Addr:                   addr,
				Name:                   _fasthttpHostClientName,
				IsTLS:                  p.opt.tlsConfig != nil,
				TLSConfig:              p.opt.tlsConfig,
				DisablePathNormalizing: p.opt.disablePathNormalizing,
			}
			p.clients = append(p.clients, client)
		}

		return nil
	}

	// not open balancer
	p.bla = nil
	addr := p.opt.addresses[0]
	client := &fasthttp.HostClient{
		Addr:                   addr,
		Name:                   _fasthttpHostClientName,
		IsTLS:                  p.opt.tlsConfig != nil,
		TLSConfig:              p.opt.tlsConfig,
		DisablePathNormalizing: p.opt.disablePathNormalizing,
		MaxConnDuration:        p.opt.maxConnDuration,
	}
	p.clients = append(p.clients, client)

	return nil
}

func (p *ReverseProxy) getClient() *fasthttp.HostClient {
	if p.clients == nil {
		// closed
		panic("ReverseProxy has been closed")
	}

	if p.bla != nil {
		// bla has been opened
		idx := p.bla.Distribute()

		return p.clients[idx]
	}

	return p.clients[0]
}

// ServeHTTP ReverseProxy to serve
// ref to: https://golang.org/src/net/http/httputil/reverseproxy.go#L169
func (p *ReverseProxy) ServeHTTP(ctx *fasthttp.RequestCtx) {
	req := &ctx.Request
	res := &ctx.Response

	// prepare request(replace headers and some URL host)
	ip, _, err := net.SplitHostPort(ctx.RemoteAddr().String())
	if err == nil {
		req.Header.Add("X-Forwarded-For", ip)
	}

	for _, h := range hopHeaders {
		req.Header.Del(h)
	}

	c := p.getClient()
	debugf(p.opt.debug, p.opt.logger, "rev request headers to proxy, addr = %s, headers = %s", c.Addr, req.Header.String())

	// assign the host to support virtual hosting, aka shared web hosting (one IP, multiple domains)
	req.SetHost(c.Addr)

	// execute the request and rev response with timeout
	err = p.doWithTimeout(c, req, res)
	if err != nil {
		errorf(p.opt.logger, "p.doWithTimeout failed, err = %v, status = %d", err, res.StatusCode())
		res.SetStatusCode(http.StatusInternalServerError)

		if errors.Is(err, fasthttp.ErrTimeout) {
			res.SetStatusCode(http.StatusRequestTimeout)
		}

		res.SetBody([]byte(err.Error()))

		return
	}

	// deal with response headers
	debugf(
		p.opt.debug, p.opt.logger,
		"rev response headers from proxy, addr = %s, headers = %s",
		c.Addr, res.Header.String(),
	)

	for _, h := range hopHeaders {
		res.Header.Del(h)
	}
}

// doWithTimeout calls fasthttp.HostClient Do or DoTimeout, this is depends on p.opt.timeout
func (p *ReverseProxy) doWithTimeout(pc *fasthttp.HostClient, req *fasthttp.Request, res *fasthttp.Response) error {
	if p.opt.timeout <= 0 {
		return pc.Do(req, res)
	}

	return pc.DoTimeout(req, res, p.opt.timeout)
}

// SetClient ...
func (p *ReverseProxy) SetClient(addr string) *ReverseProxy {
	for idx := range p.clients {
		p.clients[idx].Addr = addr
	}

	return p
}

// Reset ...
func (p *ReverseProxy) Reset() {
	for idx := range p.clients {
		p.clients[idx].Addr = ""
	}
}

// Close ... clear and release
func (p *ReverseProxy) Close() {
	p.clients = nil
	p.opt = nil
	p.bla = nil
}

// Hop-by-hop headers. These are removed when sent to the backend.
// As of RFC 7230, hop-by-hop headers are required to appear in the
// Connection header field. These are the headers defined by the
// obsoleted RFC 2616 (section 13.5.1) and are used for backward
// compatibility.
var hopHeaders = []string{
	"Connection",          // Connection
	"Proxy-Connection",    // non-standard but still sent by libcurl and rejected by e.g. google
	"Keep-Alive",          // Keep-Alive
	"Proxy-Authenticate",  // Proxy-Authenticate
	"Proxy-Authorization", // Proxy-Authorization
	"Te",                  // canonicalized version of "TE"
	"Trailer",             // not Trailers per URL above; https://www.rfc-editor.org/errata_search.php?eid=4522
	"Transfer-Encoding",   // Transfer-Encoding
	"Upgrade",             // Upgrade
}
