package rtsp

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"hash"
	"io"
	"strings"

	parser "gitee.com/raininfall/header-parser"
)

// Authorization machine of RFC2617
type Authorization interface {
	Inject(*Request) error
}

type authorizationDigest struct {
	username string
	password string
	realm    string
	nonce    string
	// request
	method string
	uri    string
	// digest
	newHash func() hash.Hash
	A1      func() string
	A2      func() string
	KD      func(string, string) string
}

var authorizationDigestDecoder = parser.BuildDividorDecoder("Authorization", ",").
	AddSub(parser.BuildDividorDecoder("realm", "=").
		AddSub(parser.BuildConstStringDecoder("name", "realm").Build()).
		AddSub(parser.BuildStringDecoder("value").Build().Wrap(parser.QuoteWith(`"`))).
		Build()).
	AddSub(parser.BuildDividorDecoder("nonce", "=").
		AddSub(parser.BuildConstStringDecoder("name", "nonce").Build()).
		AddSub(parser.BuildStringDecoder("value").Build().Wrap(parser.QuoteWith(`"`))).
		Build()).
	Build()

// NewAuthenticate returns
func NewAuthenticate(wwwAuthenticate, username, password string) (Authorization, error) {
	params := strings.SplitN(wwwAuthenticate, " ", 2)
	if len(params) != 2 {
		return nil, &badStringError{"Invalid format", wwwAuthenticate}
	}
	switch params[0] {
	case "Digest":
		return newAuthorizationDigest(username, password, params[1])
	default:
		return nil, &badStringError{"Unknown challenge", wwwAuthenticate}
	}
}

func newAuthorizationDigest(username, password, param string) (Authorization, error) {
	result, err := authorizationDigestDecoder.Do(param)
	if err != nil {
		return nil, err
	}

	authorization := &authorizationDigest{
		username: username,
		password: password,
	}
	authorization.realm, err = result.Sub("realm").Sub("value").ValueString()
	if err != nil {
		return nil, err
	}
	authorization.nonce, err = result.Sub("nonce").Sub("value").ValueString()
	if err != nil {
		return nil, err
	}

	authorization.newHash = md5.New
	authorization.A1 = authorization.A1MD5
	authorization.A2 = authorization.A2Auth
	authorization.KD = authorization.KDMD5

	return authorization, nil
}

func (a *authorizationDigest) H(s string) string {
	h := a.newHash()
	io.WriteString(h, s)
	return hex.EncodeToString(h.Sum(nil))
}

func (a *authorizationDigest) A1MD5() string {
	return fmt.Sprintf("%s:%s:%s", a.username, a.realm, a.password)
}

func (a *authorizationDigest) A2Auth() string {
	return fmt.Sprintf("%s:%s", a.method, a.uri)
}

func (a *authorizationDigest) KDMD5(secret, data string) string {
	return a.H((fmt.Sprintf("%s:%s", secret, data)))
}

func (a *authorizationDigest) digest() string {
	secret := a.H(a.A1())
	data := fmt.Sprintf("%s:%s", a.nonce, a.H(a.A2()))
	return a.KD(secret, data)
}

func (a *authorizationDigest) Inject(request *Request) error {
	a.method = request.Method
	a.uri = request.RawURL

	line := make([]string, 0, 5)
	line = append(line, fmt.Sprintf(`Digest username="%s"`, a.username))
	line = append(line, fmt.Sprintf(`realm="%s"`, a.realm))
	line = append(line, fmt.Sprintf(`nonce="%s"`, a.nonce))
	line = append(line, fmt.Sprintf(`uri="%s"`, a.uri))
	line = append(line, fmt.Sprintf(`response="%s"`, a.digest()))

	request.Header.Set("Authorization", strings.Join(line, ", "))

	return nil
}