package hammerhead

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"encoding/base32"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"net/http"
	"sort"
	"strings"
	"time"

	"golang.org/x/crypto/bcrypt"
)

const sessionTTL = 14 * 24 * time.Hour

type principal struct {
	User User    // copy; zero ID when authenticated by API key
	Key  *APIKey // set for API-key callers
}

type ticket struct {
	userID  string
	expires time.Time
}

func clientIP(r *http.Request) string {
	h, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return h
}

func (s *Server) principalFrom(r *http.Request) *principal {
	s.mu.Lock()
	defer s.mu.Unlock()
	if c, err := r.Cookie("hh_session"); err == nil {
		if sess := s.st.Sessions[sha(c.Value)]; sess != nil && time.Now().Before(sess.Expires) {
			if u := s.findUser(sess.UserID); u != nil {
				return &principal{User: *u}
			}
		}
	}
	if tok := bearer(r); strings.HasPrefix(tok, "hhk_") {
		parts := strings.SplitN(tok, "_", 3)
		if len(parts) == 3 {
			for _, k := range s.st.APIKeys {
				if k.ID == parts[1] && hmac.Equal([]byte(k.Hash), []byte(sha(parts[2]))) {
					k.LastUsed = time.Now()
					kc := *k
					if k.OrgID == "" { // a server key acts as a server administrator
						return &principal{Key: &kc, User: User{Email: "apikey:" + k.Name, ServerAdmin: true}}
					}
					return &principal{Key: &kc}
				}
			}
		}
	}
	return nil
}

func (s *Server) withUser(fn func(http.ResponseWriter, *http.Request, *principal)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		p := s.principalFrom(r)
		if p == nil || p.User.ID == "" {
			fail(w, 401, errors.New("Sign in to continue."))
			return
		}
		fn(w, r, p)
	}
}

// serverKey reports whether the caller authenticated with a server-level API key.
func (p *principal) serverKey() bool { return p.Key != nil && p.Key.OrgID == "" }

// withServerAdmin lets signed-in server administrators and server API keys through.
func (s *Server) withServerAdmin(fn func(http.ResponseWriter, *http.Request, *principal)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		p := s.principalFrom(r)
		if p == nil || (p.User.ID == "" && !p.serverKey()) {
			fail(w, 401, errors.New("Sign in to continue."))
			return
		}
		if !p.User.ServerAdmin {
			fail(w, 403, errors.New("server administrators only"))
			return
		}
		fn(w, r, p)
	}
}

// orgCtx is the caller's standing inside one org.
type orgCtx struct {
	Org   Org
	Admin bool
	User  User
	Key   *APIKey
	Actor string
}

func (s *Server) withOrg(needAdmin bool, fn func(http.ResponseWriter, *http.Request, *orgCtx)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		p := s.principalFrom(r)
		if p == nil {
			fail(w, 401, errors.New("Sign in to continue."))
			return
		}
		id := r.PathValue("org")
		s.mu.Lock()
		org := s.findOrg(id)
		c := &orgCtx{User: p.User, Key: p.Key}
		allowed := false
		if org != nil {
			c.Org = *org
			switch {
			case p.serverKey():
				allowed, c.Admin = true, true
				c.Actor = "apikey:" + p.Key.Name
			case p.Key != nil:
				allowed = p.Key.OrgID == org.ID
				c.Admin = allowed
				c.Actor = "apikey:" + p.Key.Name
			case p.User.ServerAdmin:
				allowed, c.Admin = true, true
				c.Actor = p.User.Email
			default:
				if m := s.member(p.User.ID, org.ID); m != nil {
					allowed = true
					if role := s.findRole(m.RoleID); role != nil {
						c.Admin = role.Admin
					}
				}
				c.Actor = p.User.Email
			}
		}
		s.mu.Unlock()
		if !allowed {
			fail(w, 404, errors.New("no such organization"))
			return
		}
		if needAdmin && !c.Admin {
			fail(w, 403, errors.New("organization administrators only"))
			return
		}
		fn(w, r, c)
	}
}

// ---------- first run ----------

func hashPassword(pw string) (string, error) {
	if len(pw) < 8 {
		return "", errors.New("password must be at least 8 characters")
	}
	b, err := bcrypt.GenerateFromPassword([]byte(pw), 10)
	return string(b), err
}

// newOrgLocked creates an org with its two built-in roles. Caller holds s.mu.
func (s *Server) newOrgLocked(name string, allowLocal bool) (*Org, *Role) {
	o := &Org{ID: newID(), Name: name, AllowLocal: allowLocal}
	admin := &Role{ID: newID(), OrgID: o.ID, Name: "Admin", Admin: true, Builtin: true}
	member := &Role{ID: newID(), OrgID: o.ID, Name: "Member", Builtin: true}
	s.st.Orgs = append(s.st.Orgs, o)
	s.st.Roles = append(s.st.Roles, admin, member)
	return o, admin
}

func (s *Server) createFirstAdmin(email, name, password, orgName string) error {
	hash, err := hashPassword(password)
	if err != nil {
		return err
	}
	email = strings.ToLower(strings.TrimSpace(email))
	if !strings.Contains(email, "@") {
		return errors.New("enter a valid email address")
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if len(s.st.Users) != 0 {
		return errors.New("Hammerhead is already set up. Sign in instead.")
	}
	u := &User{ID: newID(), Email: email, Name: name, PassHash: hash, ServerAdmin: true, Created: time.Now()}
	s.st.Users = append(s.st.Users, u)
	org, adminRole := s.newOrgLocked(orgName, true)
	s.st.Members = append(s.st.Members, &Member{UserID: u.ID, OrgID: org.ID, RoleID: adminRole.ID})
	return s.save()
}

func (s *Server) setup(w http.ResponseWriter, r *http.Request) {
	var in struct{ Email, Name, Password, OrgName string }
	if !decode(w, r, &in) {
		return
	}
	if strings.TrimSpace(in.OrgName) == "" {
		in.OrgName = "Default"
	}
	if strings.TrimSpace(in.Name) == "" {
		in.Name = "Admin"
	}
	if err := s.createFirstAdmin(in.Email, in.Name, in.Password, in.OrgName); err != nil {
		fail(w, 400, err)
		return
	}
	ok(w)
}

// ---------- sessions ----------

func (s *Server) secureCookies(r *http.Request) bool {
	return r.TLS != nil || strings.HasPrefix(s.st.Settings.DashboardURL, "https://")
}

func (s *Server) startSession(w http.ResponseWriter, r *http.Request, userID string) {
	tok := randHex(24)
	now := time.Now()
	s.mu.Lock()
	for k, v := range s.st.Sessions {
		if now.After(v.Expires) {
			delete(s.st.Sessions, k)
		}
	}
	s.st.Sessions[sha(tok)] = &Session{UserID: userID, Expires: now.Add(sessionTTL), Created: now, IP: clientIP(r), UA: trim(r.UserAgent(), 200)}
	secure := s.secureCookies(r)
	s.save()
	s.mu.Unlock()
	http.SetCookie(w, &http.Cookie{Name: "hh_session", Value: tok, Path: "/", HttpOnly: true,
		Secure: secure, SameSite: http.SameSiteLaxMode, MaxAge: int(sessionTTL.Seconds())})
}

func (s *Server) logout(w http.ResponseWriter, r *http.Request, _ *principal) {
	if c, err := r.Cookie("hh_session"); err == nil {
		s.mu.Lock()
		delete(s.st.Sessions, sha(c.Value))
		s.save()
		s.mu.Unlock()
	}
	http.SetCookie(w, &http.Cookie{Name: "hh_session", Value: "", Path: "/", MaxAge: -1})
	ok(w)
}

// dummyHash lets a missing user cost the same as a wrong password.
var dummyHash, _ = bcrypt.GenerateFromPassword([]byte("not-a-real-password"), 10)

func (s *Server) login(w http.ResponseWriter, r *http.Request) {
	var in struct{ Email, Password string }
	if !decode(w, r, &in) {
		return
	}
	email := strings.ToLower(strings.TrimSpace(in.Email))
	if !s.allow("login|"+clientIP(r)+"|"+email, 8, 5*time.Minute) {
		fail(w, 429, errors.New("Too many attempts. Wait a few minutes and try again."))
		return
	}
	s.mu.Lock()
	u := s.userByEmail(email)
	hash := string(dummyHash)
	if u != nil && u.PassHash != "" {
		hash = u.PassHash
	}
	s.mu.Unlock()
	good := bcrypt.CompareHashAndPassword([]byte(hash), []byte(in.Password)) == nil
	if u == nil || !good || u.PassHash == "" {
		s.audit("", email, "login.failed", clientIP(r))
		time.Sleep(300 * time.Millisecond)
		fail(w, 401, errors.New("That email and password do not match. Check them and try again."))
		return
	}
	if u.TOTPEnabled {
		t := randHex(16)
		s.mu.Lock()
		s.tickets[t] = ticket{userID: u.ID, expires: time.Now().Add(5 * time.Minute)}
		s.mu.Unlock()
		writeJSON(w, 200, map[string]any{"needTotp": true, "ticket": t})
		return
	}
	s.startSession(w, r, u.ID)
	s.audit("", u.Email, "login", clientIP(r))
	ok(w)
}

func (s *Server) loginTOTP(w http.ResponseWriter, r *http.Request) {
	var in struct{ Ticket, Code string }
	if !decode(w, r, &in) {
		return
	}
	if !s.allow("totp|"+clientIP(r), 10, 5*time.Minute) {
		fail(w, 429, errors.New("Too many attempts. Wait a few minutes and try again."))
		return
	}
	s.mu.Lock()
	t, found := s.tickets[in.Ticket]
	var u *User
	if found && time.Now().Before(t.expires) {
		u = s.findUser(t.userID)
	}
	good := false
	if u != nil {
		code := strings.ReplaceAll(strings.TrimSpace(in.Code), " ", "")
		if totpOK(u.TOTPSecret, code) {
			good = true
		} else {
			h := sha(strings.ToLower(code))
			for i, b := range u.BackupCodes {
				if b == h {
					u.BackupCodes = append(append([]string{}, u.BackupCodes[:i]...), u.BackupCodes[i+1:]...)
					good = true
					break
				}
			}
		}
	}
	if good {
		delete(s.tickets, in.Ticket)
		s.save()
	}
	s.mu.Unlock()
	if !good {
		fail(w, 401, errors.New("That code is not right. Check your authenticator and try again."))
		return
	}
	s.startSession(w, r, u.ID)
	s.audit("", u.Email, "login", clientIP(r)+" (2fa)")
	ok(w)
}

// ---------- TOTP ----------

var b32 = base32.StdEncoding.WithPadding(base32.NoPadding)

func hotp(key []byte, counter uint64) string {
	var b [8]byte
	binary.BigEndian.PutUint64(b[:], counter)
	m := hmac.New(sha1.New, key)
	m.Write(b[:])
	sum := m.Sum(nil)
	off := sum[len(sum)-1] & 0xf
	v := (binary.BigEndian.Uint32(sum[off:off+4]) & 0x7fffffff) % 1000000
	return fmt.Sprintf("%06d", v)
}

func totpOK(secretB32, code string) bool {
	key, err := b32.DecodeString(secretB32)
	if err != nil || len(code) != 6 {
		return false
	}
	now := uint64(time.Now().Unix() / 30)
	for _, d := range []int{-1, 0, 1} {
		if hmac.Equal([]byte(hotp(key, now+uint64(d))), []byte(code)) {
			return true
		}
	}
	return false
}

// ---------- /api/me ----------

type orgBrief struct {
	ID    string `json:"id"`
	Name  string `json:"name"`
	Role  string `json:"role"`
	Admin bool   `json:"admin"`
}

func (s *Server) publicInfo(w http.ResponseWriter, r *http.Request) {
	s.mu.Lock()
	defer s.mu.Unlock()
	writeJSON(w, 200, map[string]any{"needsSetup": len(s.st.Users) == 0, "providers": s.providerList()})
}

func (s *Server) providerList() []map[string]string {
	out := []map[string]string{}
	for _, p := range s.st.IdPs {
		out = append(out, map[string]string{"id": p.ID, "name": p.Name})
	}
	return out
}

func (s *Server) me(w http.ResponseWriter, r *http.Request, p *principal) {
	s.mu.Lock()
	defer s.mu.Unlock()
	u := s.findUser(p.User.ID)
	orgs := []orgBrief{}
	for _, o := range s.st.Orgs {
		if u.ServerAdmin {
			orgs = append(orgs, orgBrief{ID: o.ID, Name: o.Name, Role: "Server admin", Admin: true})
			continue
		}
		if m := s.member(u.ID, o.ID); m != nil {
			role := s.findRole(m.RoleID)
			b := orgBrief{ID: o.ID, Name: o.Name}
			if role != nil {
				b.Role, b.Admin = role.Name, role.Admin
			}
			orgs = append(orgs, b)
		}
	}
	writeJSON(w, 200, map[string]any{
		"id": u.ID, "email": u.Email, "name": u.Name, "serverAdmin": u.ServerAdmin,
		"totpEnabled": u.TOTPEnabled, "hasPassword": u.PassHash != "", "orgs": orgs,
		"backupCodes": len(u.BackupCodes), "dashboardUrl": s.st.Settings.DashboardURL,
	})
}

func (s *Server) updateMe(w http.ResponseWriter, r *http.Request, p *principal) {
	var in struct{ Name string }
	if !decode(w, r, &in) {
		return
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if n := strings.TrimSpace(in.Name); n != "" {
		s.findUser(p.User.ID).Name = n
	}
	s.save()
	ok(w)
}

func (s *Server) changePassword(w http.ResponseWriter, r *http.Request, p *principal) {
	var in struct{ Current, New string }
	if !decode(w, r, &in) {
		return
	}
	if p.User.PassHash != "" && bcrypt.CompareHashAndPassword([]byte(p.User.PassHash), []byte(in.Current)) != nil {
		fail(w, 400, errors.New("Your current password is not right. Check it and try again."))
		return
	}
	hash, err := hashPassword(in.New)
	if err != nil {
		fail(w, 400, err)
		return
	}
	s.mu.Lock()
	s.findUser(p.User.ID).PassHash = hash
	// sign out every other session
	cur := ""
	if c, err := r.Cookie("hh_session"); err == nil {
		cur = sha(c.Value)
	}
	for k, sess := range s.st.Sessions {
		if sess.UserID == p.User.ID && k != cur {
			delete(s.st.Sessions, k)
		}
	}
	s.save()
	s.mu.Unlock()
	s.audit("", p.User.Email, "password.change", "")
	ok(w)
}

func (s *Server) totpSetup(w http.ResponseWriter, r *http.Request, p *principal) {
	if p.User.TOTPEnabled {
		fail(w, 400, errors.New("two-factor is already on"))
		return
	}
	raw := make([]byte, 20)
	rand.Read(raw)
	secret := b32.EncodeToString(raw)
	s.mu.Lock()
	s.findUser(p.User.ID).TOTPSecret = secret
	s.save()
	s.mu.Unlock()
	writeJSON(w, 200, map[string]string{"secret": secret,
		"uri": fmt.Sprintf("otpauth://totp/Hammerhead:%s?secret=%s&issuer=Hammerhead", p.User.Email, secret)})
}

func (s *Server) totpEnable(w http.ResponseWriter, r *http.Request, p *principal) {
	var in struct{ Code string }
	if !decode(w, r, &in) {
		return
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	u := s.findUser(p.User.ID)
	if u.TOTPSecret == "" || !totpOK(u.TOTPSecret, strings.TrimSpace(in.Code)) {
		fail(w, 400, errors.New("That code does not match. Check that your authenticator's clock is right and try again."))
		return
	}
	codes, hashes := make([]string, 8), make([]string, 8)
	for i := range codes {
		codes[i] = randHex(5)
		hashes[i] = sha(codes[i])
	}
	u.TOTPEnabled, u.BackupCodes = true, hashes
	s.save()
	writeJSON(w, 200, map[string]any{"backupCodes": codes})
}

func (s *Server) totpDisable(w http.ResponseWriter, r *http.Request, p *principal) {
	var in struct{ Password, Code string }
	if !decode(w, r, &in) {
		return
	}
	if p.User.PassHash != "" && bcrypt.CompareHashAndPassword([]byte(p.User.PassHash), []byte(in.Password)) != nil {
		fail(w, 400, errors.New("That password is not right. Check it and try again."))
		return
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	u := s.findUser(p.User.ID)
	if !totpOK(u.TOTPSecret, strings.TrimSpace(in.Code)) {
		fail(w, 400, errors.New("That code is not right. Check your authenticator and try again."))
		return
	}
	u.TOTPEnabled, u.TOTPSecret, u.BackupCodes = false, "", nil
	s.save()
	ok(w)
}

func trim(s string, n int) string {
	if len(s) > n {
		return s[:n]
	}
	return s
}

// ---------- your signed-in devices ----------

func (s *Server) mySessions(w http.ResponseWriter, r *http.Request, p *principal) {
	cur := ""
	if c, err := r.Cookie("hh_session"); err == nil {
		cur = sha(c.Value)
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	type row struct {
		ID      string    `json:"id"`
		Created time.Time `json:"created"`
		Expires time.Time `json:"expires"`
		IP      string    `json:"ip"`
		UA      string    `json:"ua"`
		Current bool      `json:"current"`
	}
	out := []row{}
	for k, sess := range s.st.Sessions {
		if sess.UserID == p.User.ID && time.Now().Before(sess.Expires) {
			out = append(out, row{k[:16], sess.Created, sess.Expires, sess.IP, sess.UA, k == cur})
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Created.After(out[j].Created) })
	writeJSON(w, 200, out)
}

func (s *Server) endSession(w http.ResponseWriter, r *http.Request, p *principal) {
	id := r.PathValue("id")
	s.mu.Lock()
	defer s.mu.Unlock()
	if len(id) >= 16 {
		for k, sess := range s.st.Sessions {
			if sess.UserID == p.User.ID && strings.HasPrefix(k, id) {
				delete(s.st.Sessions, k)
			}
		}
	}
	s.save()
	ok(w)
}

// endOtherSessions signs the user out everywhere except the browser they are using.
func (s *Server) endOtherSessions(w http.ResponseWriter, r *http.Request, p *principal) {
	cur := ""
	if c, err := r.Cookie("hh_session"); err == nil {
		cur = sha(c.Value)
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	for k, sess := range s.st.Sessions {
		if sess.UserID == p.User.ID && k != cur {
			delete(s.st.Sessions, k)
		}
	}
	s.save()
	ok(w)
}