package onboarding

import (
	"bytes"
	"context"
	"crypto/hmac"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"time"
	"unicode/utf8"

	"data-terminal-cloud/internal/identity"
	"github.com/dchest/captcha"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/redis/go-redis/v9"
	"golang.org/x/crypto/argon2"
)

// Bound expensive password work independently of database and request concurrency.
var passwordWorkers = make(chan struct{}, 2)
var invalidPassword = fail(401, "PASSWORD_INVALID", "手机号或密码不正确")

const passwordPrefix = "$argon2id$v=19$m=65536,t=3,p=4$"

func passwordWork(ctx context.Context, fn func() string) (string, error) {
	select {
	case passwordWorkers <- struct{}{}:
		defer func() { <-passwordWorkers }()
		if err := ctx.Err(); err != nil {
			return "", err
		}
		return fn(), nil
	default:
		return "", fail(503, "LOGIN_BUSY", "登录繁忙，请稍后重试")
	}
}
func hashPassword(ctx context.Context, password string) (string, error) {
	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		return "", err
	}
	return passwordWork(ctx, func() string {
		hash := argon2.IDKey([]byte(password), salt, 3, 64*1024, 4, 32)
		return passwordPrefix + base64.RawStdEncoding.EncodeToString(salt) + "$" + base64.RawStdEncoding.EncodeToString(hash)
	})
}
func verifyPassword(ctx context.Context, encoded, password string) (bool, error) {
	// Missing credentials still perform the same password hash work.
	salt, expected := make([]byte, 16), make([]byte, 32)
	valid := false
	if strings.HasPrefix(encoded, passwordPrefix) {
		parts := strings.Split(strings.TrimPrefix(encoded, passwordPrefix), "$")
		if len(parts) == 2 {
			s, e1 := base64.RawStdEncoding.DecodeString(parts[0])
			h, e2 := base64.RawStdEncoding.DecodeString(parts[1])
			if e1 == nil && e2 == nil && len(s) == 16 && len(h) == 32 {
				salt, expected, valid = s, h, true
			}
		}
	}
	actual, err := passwordWork(ctx, func() string { return string(argon2.IDKey([]byte(password), salt, 3, 64*1024, 4, 32)) })
	return hmac.Equal([]byte(actual), expected) && valid, err
}
func validNewPassword(password string) bool {
	n := utf8.RuneCountInString(password)
	return utf8.ValidString(password) && n >= 6 && n <= 64 && len(password) <= 256
}

func (s *Service) captchaKey(sessionHash string) string {
	return "onboarding:login-captcha:" + sessionHash
}
func (s *Service) loginCaptcha(ctx context.Context, sessionHash string) (map[string]any, error) {
	return s.newCaptcha(ctx, sessionHash, func(tx pgx.Tx) error { return s.checkLoginSession(ctx, tx, sessionHash) })
}

func (s *Service) webCaptcha(ctx context.Context, w WebSession) (map[string]any, error) {
	return s.newCaptcha(ctx, w.Hash, func(tx pgx.Tx) error {
		d, err := readDevice(ctx, tx, w.Device.ID, true)
		if err != nil {
			return err
		}
		return s.checkWebActor(ctx, tx, w.Hash, d)
	})
}

func (s *Service) newCaptcha(ctx context.Context, sessionHash string, check func(pgx.Tx) error) (map[string]any, error) {
	if err := s.limit(ctx, "captcha:session:"+sessionHash, 30, 10*time.Minute); err != nil {
		return nil, err
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return nil, err
	}
	defer tx.Rollback(ctx)
	if err = check(tx); err != nil {
		return nil, err
	}
	if s.saasCaptchaDisabled(ctx) {
		return map[string]any{"mode": "none", "captcha_id": "", "expires_at_ms": 0}, tx.Commit(ctx)
	}
	id := uuid.NewString()
	policy, err := s.readCaptchaPolicy(ctx, tx)
	if err != nil {
		return nil, err
	}
	client := captchaRequest(ctx).client
	mode := policy.mode(client)
	stored := storedCaptcha{ID: id, Mode: mode, Client: client, Version: policy.Version}
	if mode == "slide" {
		out, e := s.newSlideCaptcha(ctx, sessionHash, &stored)
		if e != nil {
			return nil, e
		}
		if e = tx.Commit(ctx); e != nil {
			return nil, e
		}
		return out, nil
	}
	if mode == "tencent" {
		stored.AppID = s.captchaAppID(client)
		if stored.AppID == "" {
			return nil, captchaUnavailable()
		}
		if err = s.redis.Set(ctx, s.captchaKey(sessionHash), captchaJSON(stored), 2*time.Minute).Err(); err != nil {
			return nil, captchaUnavailable()
		}
		if err = tx.Commit(ctx); err != nil {
			return nil, err
		}
		return map[string]any{"captcha_id": id, "mode": mode, "captcha_app_id": stored.AppID, "expires_at_ms": s.now().Add(2 * time.Minute).UnixMilli()}, nil
	}
	digits := captcha.RandomDigits(6)
	answer := make([]byte, len(digits))
	for i, d := range digits {
		answer[i] = '0' + d
	}
	var png bytes.Buffer
	if _, err = captcha.NewImage(id, digits, 240, 80).WriteTo(&png); err != nil {
		return nil, err
	}
	// Only a keyed digest is stored; the response contains a raster image, never the answer.
	stored.AnswerHash = domainMAC(s.crypt.hashKey, "CAPTCHA1\n"+sessionHash+"\n"+id+"\n"+string(answer))
	value := captchaJSON(stored)
	if err = s.redis.Set(ctx, s.captchaKey(sessionHash), value, 2*time.Minute).Err(); err != nil {
		return nil, fail(503, "CAPTCHA_UNAVAILABLE", "图片验证码暂不可用，请重试")
	}
	if err = tx.Commit(ctx); err != nil {
		return nil, err
	}
	return map[string]any{"captcha_id": id, "mode": "image", "image": "data:image/png;base64," + base64.StdEncoding.EncodeToString(png.Bytes()), "expires_at_ms": s.now().Add(2 * time.Minute).UnixMilli()}, nil
}
func (s *Service) verifyLoginCaptcha(ctx context.Context, sessionHash, id, answer string, proof ...string) error {
	if s.saasCaptchaDisabled(ctx) {
		return nil
	}
	policy, err := s.readCaptchaPolicy(ctx, s.pool)
	if err != nil {
		return captchaUnavailable()
	}
	value, err := s.consumeCaptcha(ctx, sessionHash, id)
	if err != nil && !errors.Is(err, redis.Nil) {
		return fail(503, "CAPTCHA_UNAVAILABLE", "图片验证码暂不可用，请重试")
	}
	if err != nil || !uuidV4(id) {
		return captchaInvalid()
	}
	client := captchaRequest(ctx)
	var stored storedCaptcha
	if json.Unmarshal([]byte(value), &stored) == nil {
		if stored.ID != id || stored.Client != client.client || stored.Version != policy.Version || stored.Mode != policy.mode(client.client) {
			return captchaInvalid()
		}
		if stored.Mode == "slide" {
			if len(proof) != 2 || !validB64(proof[0], 32) || proof[1] != "" || answer != "" ||
				len(stored.Solution) != 0 || stored.ProofHash == "" || stored.ExpiresAt <= s.now().UnixMilli() ||
				!hmac.Equal([]byte(stored.ProofHash), []byte(s.slideProofHash(sessionHash, id, proof[0]))) {
				return captchaInvalid()
			}
			return nil
		}
		if stored.Mode == "tencent" {
			if len(proof) != 2 || proof[0] == "" || len(proof[0]) > 4096 || len(proof[1]) > 256 || strings.HasPrefix(proof[0], "trerror") {
				return captchaInvalid()
			}
			if stored.AppID != s.captchaAppID(client.client) || s.options.Captcha == nil {
				return captchaUnavailable()
			}
			// Burn a ticket before network I/O, across all sessions/realms/endpoints.
			fresh, e := s.redis.SetNX(ctx, "onboarding:captcha-ticket:"+digest([]byte(proof[0])), "used", 10*time.Minute).Result()
			if e != nil {
				return captchaUnavailable()
			}
			if !fresh {
				return captchaInvalid()
			}
			return s.options.Captcha.Verify(ctx, client.client, proof[0], proof[1], client.ip)
		}
		value = stored.ID + ":" + stored.AnswerHash
	} else if policy.mode(client.client) != "image" {
		return captchaInvalid()
	}
	expected := id + ":" + domainMAC(s.crypt.hashKey, "CAPTCHA1\n"+sessionHash+"\n"+id+"\n"+answer)
	if err != nil || !uuidV4(id) || !smsCodePattern.MatchString(answer) || !hmac.Equal([]byte(value), []byte(expected)) {
		return fail(401, "CAPTCHA_INVALID", "图片验证码错误或已过期，请重新输入")
	}
	return nil
}

func (s *Service) saasCaptchaDisabled(ctx context.Context) bool {
	return s.options.SaaSCaptchaDisabled && captchaRequest(ctx).client == "saas"
}

type PasswordLoginInput struct {
	TrustedDeviceToken string `json:"trusted_device_token,omitempty"`
	CaptchaTicket      string `json:"captcha_ticket,omitempty"`
	CaptchaRandstr     string `json:"captcha_randstr,omitempty"`
	Phone              string `json:"phone_e164"`
	Password           string `json:"password"`
	CaptchaID          string `json:"captcha_id"`
	CaptchaCode        string `json:"captcha_code"`
}

func (s *Service) passwordLogin(ctx context.Context, sessionHash string, in PasswordLoginInput) (BrowserCredentials, error) {
	return s.passwordLoginLinked(ctx, sessionHash, in, "")
}

func (s *Service) passwordLoginLinked(ctx context.Context, sessionHash string, in PasswordLoginInput, openID string) (BrowserCredentials, error) {
	var out BrowserCredentials
	if !phonePattern.MatchString(in.Phone) || len(in.Password) == 0 || len(in.Password) > 256 {
		return out, invalidPassword
	}
	if err := s.limit(ctx, "password:session:"+sessionHash, 20, 10*time.Minute); err != nil {
		return out, err
	}
	if err := s.verifyLoginCaptcha(ctx, sessionHash, in.CaptchaID, in.CaptchaCode, in.CaptchaTicket, in.CaptchaRandstr); err != nil {
		return out, err
	}
	phoneHash := domainMAC(s.crypt.hashKey, "PHONE1\n"+in.Phone)
	if err := s.limit(ctx, "password:phone:"+phoneHash, 10, 10*time.Minute); err != nil {
		return out, err
	}
	var encoded string
	err := s.pool.QueryRow(ctx, `SELECT COALESCE(password_hash,'') FROM `+s.phoneCredentialTable()+` WHERE phone_digest=$1`, phoneHash).Scan(&encoded)
	if err != nil && !errors.Is(err, pgx.ErrNoRows) {
		return out, err
	}
	ok, err := verifyPassword(ctx, encoded, in.Password)
	if err != nil {
		return out, err
	}
	if !ok {
		return out, invalidPassword
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return out, err
	}
	defer tx.Rollback(ctx)
	if err = s.checkLoginSession(ctx, tx, sessionHash); err != nil {
		return out, err
	}
	var p identity.Principal
	var current string
	err = tx.QueryRow(ctx, `SELECT issuer,subject,COALESCE(password_hash,'') FROM `+s.phoneCredentialTable()+` WHERE phone_digest=$1 FOR SHARE`, phoneHash).Scan(&p.Issuer, &p.Subject, &current)
	if errors.Is(err, pgx.ErrNoRows) || (err == nil && current != encoded) {
		return out, invalidPassword
	}
	if err != nil {
		return out, err
	}
	if err = s.requireTrustedLoginDevice(ctx, tx, phoneHash, current, in.TrustedDeviceToken); err != nil {
		return out, err
	}
	if openID != "" {
		if err = s.linkWeChatPhone(ctx, tx, openID, phoneHash); err != nil {
			return out, err
		}
	}
	out, err = s.attachCustomerLogin(ctx, tx, sessionHash, p, false)
	if err != nil {
		return out, err
	}
	_, err = tx.Exec(ctx, `UPDATE `+s.accountSessionTable()+` SET phone_digest=$2 WHERE token_hash=$1`, digest([]byte(out.AccountToken)), phoneHash)
	if err != nil {
		return out, err
	}
	if err = tx.Commit(ctx); err != nil {
		return out, err
	}
	return out, nil
}

type SetPasswordInput struct {
	Password        string `json:"password"`
	CurrentPassword string `json:"current_password,omitempty"`
}

func (s *Service) accountPasswordHTTP(r *http.Request) (any, error) {
	if r.URL.RawQuery != "" || r.Header.Get("Sec-Fetch-Site") == "cross-site" || (r.Method != http.MethodGet && !s.sameOrigin(r)) {
		return nil, errAccess
	}
	cookie, err := r.Cookie(s.cookieName(true))
	if err != nil {
		return nil, fail(401, "LOGIN_REQUIRED", "请先登录")
	}
	p, err := s.AccountPrincipal(r.Context(), cookie.Value, r.Header.Get("X-CSRF-Token"), r.Method != http.MethodGet)
	if err != nil {
		return nil, err
	}
	ctx := r.Context()
	sessionHash := digest([]byte(cookie.Value))
	if r.Method == http.MethodGet {
		var masked string
		var configured, recent bool
		err = s.pool.QueryRow(ctx, `SELECT c.phone_masked,c.password_hash IS NOT NULL,COALESCE(ss.phone_verified_at>$2::timestamptz-interval '10 minutes',false) FROM `+s.accountSessionTable()+` ss JOIN `+s.phoneCredentialTable()+` c ON c.phone_digest=ss.phone_digest WHERE ss.token_hash=$1`, sessionHash, s.now()).Scan(&masked, &configured, &recent)
		if errors.Is(err, pgx.ErrNoRows) {
			return map[string]any{"available": false}, nil
		}
		if err != nil {
			return nil, err
		}
		return map[string]any{"available": true, "phone_masked": masked, "password_configured": configured, "sms_verified_recently": recent}, nil
	}
	raw, err := readBody(r)
	if err != nil {
		return nil, err
	}
	var in SetPasswordInput
	if err = StrictJSON(raw, &in); err != nil {
		return nil, err
	}
	if !validNewPassword(in.Password) || len(in.CurrentPassword) > 256 {
		return nil, fail(400, "PASSWORD_POLICY", "密码须为 6–64 个字符")
	}
	if err = s.limit(ctx, "password:set:"+p.Issuer+":"+p.Subject, 10, 10*time.Minute); err != nil {
		return nil, err
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return nil, err
	}
	defer tx.Rollback(ctx)
	var phoneHash, encoded string
	var recent bool
	// Credential locking serializes changes with login, preventing a login using a concurrently reset password.
	err = tx.QueryRow(ctx, `SELECT phone_digest,COALESCE(password_hash,'') FROM `+s.phoneCredentialTable()+` WHERE issuer=$1 AND subject=$2 FOR UPDATE`, p.Issuer, p.Subject).Scan(&phoneHash, &encoded)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, fail(403, "SMS_REAUTH_REQUIRED", "请先使用短信验证码登录")
	}
	if err != nil {
		return nil, err
	}
	err = tx.QueryRow(ctx, `SELECT COALESCE(phone_verified_at>$2::timestamptz-interval '10 minutes',false) FROM `+s.accountSessionTable()+` WHERE token_hash=$1 AND phone_digest=$3 AND revoked_at IS NULL AND expires_at>now() FOR UPDATE`, sessionHash, s.now(), phoneHash).Scan(&recent)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, fail(401, "LOGIN_REQUIRED", "请重新登录")
	}
	if err != nil {
		return nil, err
	}
	if !recent {
		if in.CurrentPassword == "" {
			return nil, fail(403, "SMS_REAUTH_REQUIRED", "请使用短信验证码登录后设置密码，或输入当前密码")
		}
		ok, e := verifyPassword(ctx, encoded, in.CurrentPassword)
		if e != nil {
			return nil, e
		}
		if !ok {
			return nil, fail(401, "CURRENT_PASSWORD_INVALID", "当前密码不正确")
		}
	}
	hash, err := hashPassword(ctx, in.Password)
	if err != nil {
		return nil, err
	}
	_, err = tx.Exec(ctx, `UPDATE `+s.phoneCredentialTable()+` SET password_hash=$2,updated_at=now() WHERE phone_digest=$1`, phoneHash, hash)
	if err != nil {
		return nil, err
	}
	if err = s.updateTrustedDevicesAfterPassword(ctx, tx, phoneHash, encoded, hash, s.trustedLoginCookie(r)); err != nil {
		return nil, err
	}
	_, err = tx.Exec(ctx, `UPDATE `+s.accountSessionTable()+` ss SET phone_verified_at=NULL,revoked_at=CASE WHEN ss.token_hash=$3 THEN ss.revoked_at ELSE COALESCE(ss.revoked_at,now()) END FROM user_account ua WHERE ss.user_id=ua.id AND ua.oidc_issuer=$1 AND ua.oidc_subject=$2`, p.Issuer, p.Subject, sessionHash)
	if err != nil {
		return nil, err
	}
	if s.operatorLogin {
		return map[string]any{"password_configured": true}, tx.Commit(ctx)
	}
	setupHash := ""
	if c, e := r.Cookie(s.cookieName(false)); e == nil {
		setupHash = digest([]byte(c.Value))
	}
	_, err = tx.Exec(ctx, `UPDATE onboarding_web_session SET revoked_at=COALESCE(revoked_at,now()) WHERE issuer=$1 AND subject=$2 AND token_hash<>$3`, p.Issuer, p.Subject, setupHash)
	if err != nil {
		return nil, err
	}
	return map[string]any{"password_configured": true}, tx.Commit(ctx)
}