package server

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/gorilla/csrf"
	"github.com/gorilla/sessions"
	"github.com/tstranex/u2f"
	"go.opencensus.io/trace"

	"git.autistici.org/id/auth"
	authclient "git.autistici.org/id/auth/client"
	"git.autistici.org/id/go-sso/httputil"
	"git.autistici.org/id/go-sso/server/device"
)

const loginSessionKey = "_login"

type loginSession struct {
	*httputil.ExpiringSession

	State loginState

	// Post-login redirection URL.
	Redir string

	// Cached from the first form.
	Username string
	Password string

	// Cached from the first auth.Response.
	U2FSignRequest *u2f.WebSignRequest

	// Never actually serialized, just passed on
	// loginStateSuccess.
	UserInfo *auth.UserInfo
}

var defaultLoginSessionLifetime = 300 * time.Second

func newLoginSession() *loginSession {
	return &loginSession{
		ExpiringSession: httputil.NewExpiringSession(defaultLoginSessionLifetime),
		State:           loginStatePassword,
	}
}

type loginState int

const (
	loginStateNone = iota
	loginStatePassword
	loginStateOTP
	loginStateU2F
	loginStateSuccess
)

func init() {
	gob.Register(&loginSession{})
}

type loginCallbackFunc func(http.ResponseWriter, *http.Request, string, string, *auth.UserInfo) error

type loginHandler struct {
	authClient         authclient.Client
	authService        string
	u2fAppID           string
	urlPrefix          string
	devMgr             *device.Manager
	loginCallback      loginCallbackFunc
	loginSessionStore  sessions.Store
	tpl                *template.Template
	accountRecoveryURL string
}

// NewLoginHandler will wrap an http.Handler with the login workflow,
// invoking it only on successful login.
func newLoginHandler(okHandler loginCallbackFunc, devMgr *device.Manager, authClient authclient.Client, authService, u2fAppID, urlPrefix, accountRecoveryURL string, tpl *template.Template, keyPairs ...[]byte) *loginHandler {
	store := sessions.NewCookieStore(keyPairs...)
	store.Options = &sessions.Options{
		HttpOnly: true,
		Secure:   true,
		MaxAge:   0,
	}
	return &loginHandler{
		authClient:         authClient,
		authService:        authService,
		u2fAppID:           u2fAppID,
		urlPrefix:          strings.TrimRight(urlPrefix, "/"),
		devMgr:             devMgr,
		loginCallback:      okHandler,
		loginSessionStore:  store,
		accountRecoveryURL: accountRecoveryURL,
		tpl:                parseEmbeddedTemplates(),
	}
}

// The login session controls the flow of the client - it's just a way
// to ensure that every step is authorized as part of the login
// sequence.
func (l *loginHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// Either fetch the current session or create a new blank one.
	httpSession, err := l.loginSessionStore.Get(req, loginSessionKey)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	session, ok := httpSession.Values["ls"].(*loginSession)
	if !ok || session == nil || !session.Valid() {
		session = newLoginSession()
		httpSession.Values["ls"] = session
	}

	// Dispatch the current state to its handler. Handlers will
	// handle the current request and either 1) validate the
	// request successfully and move to the next state, or 2)
	// return a response to the user. Handlers fall through to the
	// next state on success.
	for {
		newState, body, err := l.dispatch(w, req, session)

		// Uncaught errors result in 500s.
		if err != nil {
			log.Printf("login error: %v", err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		switch newState {
		case loginStateSuccess:
			// Successful login. Delete the login session and invoke
			// the login callback, before redirecting to the
			// original URL.
			httpSession.Options.MaxAge = -1
			if err := httpSession.Save(req, w); err != nil {
				log.Printf("login error saving session: %v", err)
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			if err := l.loginCallback(w, req, session.Username, session.Password, session.UserInfo); err != nil {
				log.Printf("login callback error: %v: user=%s", err, session.Username)
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			http.Redirect(w, req, session.Redir, http.StatusFound)
			return

		case loginStateNone:
			if err := httpSession.Save(req, w); err != nil {
				log.Printf("error saving login session: %v", err)
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			w.Write(body) // nolint
			return

		default:
			// Fall through to the next handler.
			session.State = newState
		}
	}
}

func (l *loginHandler) dispatch(w http.ResponseWriter, req *http.Request, session *loginSession) (loginState, []byte, error) {
	switch session.State {
	case loginStatePassword:
		return l.handlePassword(w, req, session)
	case loginStateOTP:
		return l.handleOTP(w, req, session)
	case loginStateU2F:
		return l.handleU2F(w, req, session)
	}
	return loginStateNone, nil, errors.New("unreachable")
}

// Handle password-based login.
func (l *loginHandler) handlePassword(w http.ResponseWriter, req *http.Request, session *loginSession) (loginState, []byte, error) {
	username := req.FormValue("username")
	password := req.FormValue("password")

	if req.Method == "GET" && session.Redir == "" {
		session.Redir = req.FormValue("r")
		// Enforce relative redirect URL (no host specified).
		if session.Redir == "" || !strings.HasPrefix(session.Redir, "/") {
			return loginStateNone, nil, errors.New("bad request")
		}
	}

	// If the request is a POST, attempt login with username/password.
	env := map[string]interface{}{
		"Error":    false,
		"Username": username,
	}
	if req.Method == "POST" && username != "" && password != "" {
		resp, err := l.makeAuthRequest(w, req, username, password, "", nil)
		if err != nil {
			return loginStateNone, nil, err
		}
		// Save username / password for later in case of
		// successful (or partially succesful) result.
		switch resp.Status {
		case auth.StatusOK:
			session.Username = username
			session.Password = password
			session.UserInfo = resp.UserInfo
			return loginStateSuccess, nil, nil
		case auth.StatusInsufficientCredentials:
			session.Username = username
			session.Password = password
			// If there is a U2F challenge in the auth
			// response, store it so we can render it
			// later.
			var nextState loginState = loginStateOTP
			if resp.TFAMethod == auth.TFAMethodU2F {
				nextState = loginStateU2F
				session.U2FSignRequest = resp.U2FSignRequest
			}
			return nextState, nil, nil
		}
		env["Error"] = true
	}

	return l.executeTemplateToBuffer(req, "login_password.html", env)
}

// Handle login with password and TOTP.
func (l *loginHandler) handleOTP(w http.ResponseWriter, req *http.Request, session *loginSession) (loginState, []byte, error) {
	otp := req.FormValue("otp")

	env := map[string]interface{}{"Error": false}
	if req.Method == "POST" && otp != "" {
		resp, err := l.makeAuthRequest(w, req, session.Username, session.Password, otp, nil)
		if err != nil {
			return loginStateNone, nil, err
		}
		if resp.Status == auth.StatusOK {
			session.UserInfo = resp.UserInfo
			return loginStateSuccess, nil, nil
		}
		env["Error"] = true
	}

	return l.executeTemplateToBuffer(req, "login_otp.html", env)
}

// Handle login with password and hardware token.
func (l *loginHandler) handleU2F(w http.ResponseWriter, req *http.Request, session *loginSession) (loginState, []byte, error) {
	u2fresponse := req.FormValue("u2f_response")

	env := map[string]interface{}{
		"U2FSignRequest": session.U2FSignRequest,
		"Error":          false,
	}
	if req.Method == "POST" && u2fresponse != "" {
		var usr u2f.SignResponse
		if err := json.Unmarshal([]byte(u2fresponse), &usr); err != nil {
			return loginStateNone, nil, err
		}

		resp, err := l.makeAuthRequest(w, req, session.Username, session.Password, "", &usr)
		if err != nil {
			return loginStateNone, nil, err
		}
		if resp.Status == auth.StatusOK {
			session.UserInfo = resp.UserInfo
			return loginStateSuccess, nil, nil
		}
		env["Error"] = true
	}

	return l.executeTemplateToBuffer(req, "login_u2f.html", env)
}

// Make the auth request to the authentication server.
func (l *loginHandler) makeAuthRequest(w http.ResponseWriter, req *http.Request, username, password, otp string, u2fResponse *u2f.SignResponse) (*auth.Response, error) {
	appID := l.u2fAppID
	if appID == "" {
		appID = u2fAppIDFromRequest(req)
	}
	ar := auth.Request{
		Service:     l.authService,
		Username:    username,
		Password:    []byte(password),
		OTP:         otp,
		DeviceInfo:  l.devMgr.GetDeviceInfoFromRequest(w, req),
		U2FResponse: u2fResponse,
		U2FAppID:    appID,
	}

	// Trace the authentication request.
	ctx, span := trace.StartSpan(req.Context(), "auth",
		trace.WithSpanKind(trace.SpanKindClient))
	span.AddAttributes(
		trace.StringAttribute("auth.user", username),
		trace.StringAttribute("auth.service", l.authService),
		trace.BoolAttribute("auth.with_password", len(password) > 0),
		trace.BoolAttribute("auth.with_otp", otp != ""),
		trace.BoolAttribute("auth.with_u2f", u2fResponse != nil),
	)
	defer span.End()

	resp, err := l.authClient.Authenticate(ctx, &ar)

	// Record the authentication response status in the trace.
	if err != nil {
		span.SetStatus(trace.Status{Code: trace.StatusCodeUnknown, Message: err.Error()})
	} else if resp.Status == auth.StatusOK {
		span.SetStatus(trace.Status{Code: trace.StatusCodeOK, Message: "OK"})
	} else {
		span.SetStatus(trace.Status{Code: trace.StatusCodePermissionDenied, Message: resp.Status.String()})
	}

	return resp, err
}

// Return a (relative) URL that will redirect the user to the login
// page and set the continue token to the original requested URL.
func (l *loginHandler) makeLoginURL(req *http.Request) string {
	v := make(url.Values)
	v.Set("r", req.URL.Path+"?"+req.URL.RawQuery)
	return fmt.Sprintf("%s/login?%s", l.urlPrefix, v.Encode())
}

// Renders a template to a buffer, with a return value that makes it
// convenient to use in login handlers.
func (l *loginHandler) executeTemplateToBuffer(req *http.Request, templateName string, data map[string]interface{}) (loginState, []byte, error) {
	data["CSRFField"] = csrf.TemplateField(req)
	data["URLPrefix"] = l.urlPrefix
	data["AccountRecoveryURL"] = l.accountRecoveryURL
	var buf bytes.Buffer
	if err := l.tpl.ExecuteTemplate(&buf, templateName, data); err != nil {
		return loginStateNone, nil, err
	}
	return loginStateNone, buf.Bytes(), nil
}

// Template helper function that encodes its input as JSON.
func toJSON(obj interface{}) string {
	data, err := json.Marshal(obj)
	if err != nil {
		return ""
	}
	return string(data)
}

// Guess the correct U2F AppID from the HTTP request.
func u2fAppIDFromRequest(r *http.Request) string {
	return fmt.Sprintf("https://%s", r.Host)
}