package main

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

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

const (
	sessionCookie   = "puppy_session"
	sessionValidity = 30 * 24 * time.Hour
)

// ctxKey is an unexported type so our context values can't collide with any
// set elsewhere.
type ctxKey int

const userIDKey ctxKey = 0

// User is the public shape returned to clients — never the password hash.
type User struct {
	ID    string `json:"id"`
	Email string `json:"email"`
}

// Auth owns everything account-related: the users/sessions tables, the shared
// invite code required to register, and whether session cookies are marked
// Secure (on behind TLS/a proxy). photosDir is needed so the first account can
// adopt legacy flat-layout photos.
type Auth struct {
	db         *sql.DB
	inviteCode string
	secure     bool
	photosDir  string
}

func newAuth(db *sql.DB, inviteCode string, secure bool, photosDir string) *Auth {
	return &Auth{db: db, inviteCode: inviteCode, secure: secure, photosDir: photosDir}
}

// ---------- users & sessions ----------

func newID() string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		panic(err) // crypto/rand failing is unrecoverable
	}
	// RFC-4122-ish v4 layout; good enough as an opaque unique id.
	b[6] = (b[6] & 0x0f) | 0x40
	b[8] = (b[8] & 0x3f) | 0x80
	return hex.EncodeToString(b[0:4]) + "-" + hex.EncodeToString(b[4:6]) + "-" +
		hex.EncodeToString(b[6:8]) + "-" + hex.EncodeToString(b[8:10]) + "-" +
		hex.EncodeToString(b[10:16])
}

// hashToken stores only the hash of a session token, so a leaked database can't
// be used to impersonate live sessions.
func hashToken(raw string) string {
	sum := sha256.Sum256([]byte(raw))
	return hex.EncodeToString(sum[:])
}

func (a *Auth) userCount() (int, error) {
	var n int
	err := a.db.QueryRow(`SELECT COUNT(*) FROM users`).Scan(&n)
	return n, err
}

var errEmailTaken = errors.New("email already registered")

func (a *Auth) createUser(email, password string) (User, error) {
	hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return User{}, err
	}
	u := User{ID: newID(), Email: email}
	_, err = a.db.Exec(
		`INSERT INTO users (id, email, password, created) VALUES (?, ?, ?, ?)`,
		u.ID, email, string(hash), time.Now().UnixMilli())
	if err != nil {
		if strings.Contains(err.Error(), "UNIQUE") {
			return User{}, errEmailTaken
		}
		return User{}, err
	}
	return u, nil
}

// verify returns the user for the given credentials, or ok=false if the email
// is unknown or the password is wrong (indistinguishable to the caller).
func (a *Auth) verify(email, password string) (User, bool) {
	var u User
	var hash string
	err := a.db.QueryRow(
		`SELECT id, email, password FROM users WHERE email = ? COLLATE NOCASE`, email,
	).Scan(&u.ID, &u.Email, &hash)
	if err != nil {
		if !errors.Is(err, sql.ErrNoRows) {
			log.Printf("verify: %v", err)
		}
		return User{}, false
	}
	if bcrypt.CompareHashAndPassword([]byte(hash), []byte(password)) != nil {
		return User{}, false
	}
	return u, true
}

// startSession mints a token, stores its hash, and returns the raw token for
// the cookie.
func (a *Auth) startSession(userID string) (string, error) {
	raw := make([]byte, 32)
	if _, err := rand.Read(raw); err != nil {
		return "", err
	}
	token := hex.EncodeToString(raw)
	now := time.Now()
	_, err := a.db.Exec(
		`INSERT INTO sessions (token, user_id, created, expires) VALUES (?, ?, ?, ?)`,
		hashToken(token), userID, now.UnixMilli(), now.Add(sessionValidity).UnixMilli())
	if err != nil {
		return "", err
	}
	return token, nil
}

// userForToken resolves a raw cookie token to a user id, honouring expiry.
func (a *Auth) userForToken(token string) (string, bool) {
	if token == "" {
		return "", false
	}
	var userID string
	var expires int64
	err := a.db.QueryRow(
		`SELECT user_id, expires FROM sessions WHERE token = ?`, hashToken(token),
	).Scan(&userID, &expires)
	if err != nil || time.Now().UnixMilli() > expires {
		return "", false
	}
	return userID, true
}

func (a *Auth) endSession(token string) {
	if token == "" {
		return
	}
	if _, err := a.db.Exec(`DELETE FROM sessions WHERE token = ?`, hashToken(token)); err != nil {
		log.Printf("endSession: %v", err)
	}
}

// adopt gives every ownerless row (legacy single-tenant data) to userID, and
// moves legacy flat-layout photos into that user's photo directory. Called once,
// when the very first account registers.
func (a *Auth) adopt(userID string) error {
	if _, err := a.db.Exec(`UPDATE events SET user_id = ? WHERE user_id = ''`, userID); err != nil {
		return err
	}
	if _, err := a.db.Exec(`UPDATE config SET user_id = ? WHERE user_id = ''`, userID); err != nil {
		return err
	}
	return a.adoptPhotos(userID)
}

// adoptPhotos moves any *.jpg sitting directly in photosDir (the pre-accounts
// flat layout) into photosDir/<userID>/.
func (a *Auth) adoptPhotos(userID string) error {
	entries, err := os.ReadDir(a.photosDir)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil
		}
		return err
	}
	dstDir := filepath.Join(a.photosDir, userID)
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(e.Name(), ".jpg") {
			continue
		}
		if err := os.MkdirAll(dstDir, 0o755); err != nil {
			return err
		}
		if err := os.Rename(
			filepath.Join(a.photosDir, e.Name()),
			filepath.Join(dstDir, e.Name()),
		); err != nil {
			return err
		}
	}
	return nil
}

// ---------- cookies & middleware ----------

func (a *Auth) setCookie(w http.ResponseWriter, token string) {
	http.SetCookie(w, &http.Cookie{
		Name:     sessionCookie,
		Value:    token,
		Path:     "/",
		HttpOnly: true,
		Secure:   a.secure,
		SameSite: http.SameSiteLaxMode,
		Expires:  time.Now().Add(sessionValidity),
	})
}

func (a *Auth) clearCookie(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     sessionCookie,
		Value:    "",
		Path:     "/",
		HttpOnly: true,
		Secure:   a.secure,
		SameSite: http.SameSiteLaxMode,
		MaxAge:   -1,
	})
}

func cookieToken(r *http.Request) string {
	c, err := r.Cookie(sessionCookie)
	if err != nil {
		return ""
	}
	return c.Value
}

// requireUser wraps a handler so it only runs for an authenticated request,
// stashing the user id in the context. Unauthenticated calls get a 401 that the
// client uses as its cue to show the login screen.
func (a *Auth) requireUser(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		userID, ok := a.userForToken(cookieToken(r))
		if !ok {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		next(w, r.WithContext(context.WithValue(r.Context(), userIDKey, userID)))
	}
}

// userID returns the authenticated user's id; only valid inside a requireUser
// handler.
func userID(r *http.Request) string {
	id, _ := r.Context().Value(userIDKey).(string)
	return id
}

// ---------- handlers ----------

type credentials struct {
	Email    string `json:"email"`
	Password string `json:"password"`
	Invite   string `json:"invite"`
}

func writeUser(w http.ResponseWriter, u User) {
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Cache-Control", "no-store")
	_ = json.NewEncoder(w).Encode(u)
}

func (a *Auth) handleRegister(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	if a.inviteCode == "" {
		http.Error(w, "registration disabled", http.StatusForbidden)
		return
	}
	var c credentials
	if err := json.NewDecoder(io.LimitReader(r.Body, 1<<16)).Decode(&c); err != nil {
		http.Error(w, "bad json", http.StatusBadRequest)
		return
	}
	// Constant-time compare so a wrong invite code can't be timed out.
	if subtle.ConstantTimeCompare([]byte(c.Invite), []byte(a.inviteCode)) != 1 {
		http.Error(w, "invalid invite code", http.StatusForbidden)
		return
	}
	email := strings.TrimSpace(strings.ToLower(c.Email))
	if !strings.Contains(email, "@") || len(email) > 200 {
		http.Error(w, "invalid email", http.StatusBadRequest)
		return
	}
	if len(c.Password) < 8 || len(c.Password) > 200 {
		http.Error(w, "password must be at least 8 characters", http.StatusBadRequest)
		return
	}

	// Whether this is the first account decides adoption of legacy data. Check
	// before insert; the users table has no other writer during registration.
	first, err := a.userCount()
	if err != nil {
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}
	u, err := a.createUser(email, c.Password)
	if errors.Is(err, errEmailTaken) {
		http.Error(w, "email already registered", http.StatusConflict)
		return
	}
	if err != nil {
		log.Printf("register: %v", err)
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}
	if first == 0 {
		if err := a.adopt(u.ID); err != nil {
			log.Printf("adopt legacy data: %v", err)
			// Non-fatal: the account exists; legacy data just stays ownerless.
		}
	}
	a.issue(w, u)
}

func (a *Auth) handleLogin(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	var c credentials
	if err := json.NewDecoder(io.LimitReader(r.Body, 1<<16)).Decode(&c); err != nil {
		http.Error(w, "bad json", http.StatusBadRequest)
		return
	}
	u, ok := a.verify(strings.TrimSpace(strings.ToLower(c.Email)), c.Password)
	if !ok {
		http.Error(w, "invalid email or password", http.StatusUnauthorized)
		return
	}
	a.issue(w, u)
}

// issue starts a session, sets the cookie, and returns the user.
func (a *Auth) issue(w http.ResponseWriter, u User) {
	token, err := a.startSession(u.ID)
	if err != nil {
		log.Printf("start session: %v", err)
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}
	a.setCookie(w, token)
	writeUser(w, u)
}

func (a *Auth) handleLogout(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	a.endSession(cookieToken(r))
	a.clearCookie(w)
	w.WriteHeader(http.StatusNoContent)
}

// handleMe reports the current account. Wrapped in requireUser, so reaching it
// means the session is valid.
func (a *Auth) handleMe(w http.ResponseWriter, r *http.Request) {
	var u User
	err := a.db.QueryRow(
		`SELECT id, email FROM users WHERE id = ?`, userID(r),
	).Scan(&u.ID, &u.Email)
	if err != nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	writeUser(w, u)
}