package main

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/sha256"
	"database/sql"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"log"
	"os"
	"regexp"
	"strings"

	"git.118im.com/im_team/olons_server/internal/config"
	_ "github.com/go-sql-driver/mysql"
)

var phonePattern = regexp.MustCompile(`^1[3-9][0-9]{9}$`)

type memberPhone struct {
	id         int64
	tenantID   int64
	ciphertext string
	hash       string
	mask       string
	phone      string
}

func main() {
	cfg, err := config.Load()
	if err != nil {
		log.Fatal(err)
	}
	db, err := sql.Open("mysql", strings.TrimPrefix(cfg.MySQL.DSN, "mysql://"))
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	ctx := context.Background()
	hasMobile, err := columnExists(ctx, db, "mobile")
	if err != nil {
		log.Fatal(err)
	}
	hasCipher, err := columnExists(ctx, db, "mobile_cipher")
	if err != nil {
		log.Fatal(err)
	}
	if hasMobile && !hasCipher {
		log.Print("hg_app_member already uses the mobile column")
		return
	}
	if !hasCipher {
		log.Fatal("mobile_cipher is missing before migration")
	}

	aeads := newAEADs(append(splitKeys(os.Getenv("OLONS_SERVER_MEMBER_MIGRATION_DATA_KEYS")), cfg.Auth.DataEncryptionKey))
	lookupKeys := decodeKeys(append(splitKeys(os.Getenv("OLONS_SERVER_MEMBER_MIGRATION_LOOKUP_KEYS")), cfg.Auth.LookupHMACKey))
	overrides, err := parseOverrides(os.Getenv("OLONS_SERVER_MEMBER_MIGRATION_OVERRIDES"))
	if err != nil {
		log.Fatal(err)
	}
	if len(aeads) == 0 && len(lookupKeys) == 0 {
		log.Fatal("no valid historical phone keys were provided")
	}
	phones, err := loadPhones(ctx, db, aeads, lookupKeys, overrides)
	if err != nil {
		log.Fatal(err)
	}
	if err := validateUnique(phones); err != nil {
		log.Fatal(err)
	}
	if !hasMobile {
		if _, err := db.ExecContext(ctx, `ALTER TABLE hg_app_member ADD COLUMN mobile varchar(20) NULL AFTER username`); err != nil {
			log.Fatal(err)
		}
	}
	if err := updatePhones(ctx, db, phones); err != nil {
		log.Fatal(err)
	}
	if _, err := db.ExecContext(ctx, `ALTER TABLE hg_app_member
		MODIFY COLUMN mobile varchar(20) NOT NULL,
		DROP INDEX uk_app_member_mobile_hash,
		ADD UNIQUE KEY uk_app_member_mobile (tenant_id,mobile),
		DROP COLUMN mobile_cipher,
		DROP COLUMN mobile_hash,
		DROP COLUMN mobile_mask`); err != nil {
		log.Fatal(err)
	}
	log.Printf("migrated %d member phone numbers", len(phones))
}

func columnExists(ctx context.Context, db *sql.DB, column string) (bool, error) {
	var count int
	err := db.QueryRowContext(ctx, `SELECT COUNT(*) FROM information_schema.COLUMNS
		WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME='hg_app_member' AND COLUMN_NAME=?`, column).Scan(&count)
	return count == 1, err
}

func loadPhones(ctx context.Context, db *sql.DB, aeads []cipher.AEAD, lookupKeys [][]byte, overrides map[int64]string) ([]memberPhone, error) {
	rows, err := db.QueryContext(ctx, `SELECT id,tenant_id,mobile_cipher,mobile_hash,mobile_mask FROM hg_app_member ORDER BY id`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	phones := make([]memberPhone, 0)
	failed := make([]string, 0)
	for rows.Next() {
		var item memberPhone
		if err := rows.Scan(&item.id, &item.tenantID, &item.ciphertext, &item.hash, &item.mask); err != nil {
			return nil, err
		}
		item.phone = recoverPhone(item, aeads, lookupKeys)
		if item.phone == "" {
			item.phone = overrides[item.id]
		}
		if !phonePattern.MatchString(item.phone) {
			failed = append(failed, fmt.Sprintf("%d", item.id))
		} else if len(item.mask) == 11 && item.phone[:3]+"****"+item.phone[7:] != item.mask {
			return nil, fmt.Errorf("member %d override does not match its phone mask", item.id)
		}
		phones = append(phones, item)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	if len(failed) > 0 {
		return nil, fmt.Errorf("phone recovery failed for member IDs: %s", strings.Join(failed, ","))
	}
	return phones, nil
}

func validateUnique(phones []memberPhone) error {
	seen := make(map[string]int64, len(phones))
	for _, item := range phones {
		key := fmt.Sprintf("%d:%s", item.tenantID, item.phone)
		if otherID, ok := seen[key]; ok {
			return fmt.Errorf("members %d and %d have duplicate phone numbers", otherID, item.id)
		}
		seen[key] = item.id
	}
	return nil
}

func updatePhones(ctx context.Context, db *sql.DB, phones []memberPhone) error {
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()
	for _, item := range phones {
		result, err := tx.ExecContext(ctx, `UPDATE hg_app_member SET mobile=? WHERE id=? AND tenant_id=?`, item.phone, item.id, item.tenantID)
		if err != nil {
			return err
		}
		affected, err := result.RowsAffected()
		if err != nil {
			return err
		}
		if affected != 1 {
			return fmt.Errorf("member %d update affected %d rows", item.id, affected)
		}
	}
	return tx.Commit()
}

func splitKeys(value string) []string {
	return strings.FieldsFunc(value, func(r rune) bool { return r == ',' || r == '\n' })
}

func parseOverrides(value string) (map[int64]string, error) {
	result := make(map[int64]string)
	for _, pair := range splitKeys(value) {
		var id int64
		var phone string
		if _, err := fmt.Sscanf(pair, "%d=%s", &id, &phone); err != nil || id <= 0 || !phonePattern.MatchString(phone) {
			return nil, fmt.Errorf("invalid member phone override")
		}
		result[id] = phone
	}
	return result, nil
}

func newAEADs(encodedKeys []string) []cipher.AEAD {
	result := make([]cipher.AEAD, 0, len(encodedKeys))
	for _, encodedKey := range encodedKeys {
		key, err := decodeKey(encodedKey)
		if err != nil || len(key) != 32 {
			continue
		}
		block, err := aes.NewCipher(key)
		if err != nil {
			continue
		}
		aead, err := cipher.NewGCM(block)
		if err == nil {
			result = append(result, aead)
		}
	}
	return result
}

func decodeKeys(encodedKeys []string) [][]byte {
	result := make([][]byte, 0, len(encodedKeys))
	for _, encodedKey := range encodedKeys {
		key, err := decodeKey(encodedKey)
		if err == nil && len(key) >= 32 {
			result = append(result, key)
		}
	}
	return result
}

func decodeKey(value string) ([]byte, error) {
	value = strings.TrimSpace(value)
	if decoded, err := base64.StdEncoding.DecodeString(value); err == nil {
		return decoded, nil
	}
	return base64.RawStdEncoding.DecodeString(value)
}

func recoverPhone(item memberPhone, aeads []cipher.AEAD, lookupKeys [][]byte) string {
	for _, aead := range aeads {
		for _, encoding := range []*base64.Encoding{base64.RawStdEncoding, base64.RawURLEncoding} {
			plain, err := decrypt(aead, encoding, item.ciphertext)
			phone := strings.TrimPrefix(plain, "+86")
			if err == nil && phonePattern.MatchString(phone) {
				return phone
			}
		}
	}
	if len(item.mask) != 11 || item.mask[3:7] != "****" {
		return ""
	}
	for middle := 0; middle < 10000; middle++ {
		phone := fmt.Sprintf("%s%04d%s", item.mask[:3], middle, item.mask[7:])
		for _, value := range []string{phone, "+86" + phone} {
			digest := sha256.Sum256([]byte(value))
			if hmac.Equal([]byte(strings.ToLower(item.hash)), []byte(hex.EncodeToString(digest[:]))) {
				return phone
			}
		}
		for _, key := range lookupKeys {
			for _, value := range []string{phone, "+86" + phone} {
				mac := hmac.New(sha256.New, key)
				_, _ = mac.Write([]byte(value))
				if hmac.Equal([]byte(strings.ToLower(item.hash)), []byte(hex.EncodeToString(mac.Sum(nil)))) {
					return phone
				}
			}
		}
	}
	return ""
}

func decrypt(aead cipher.AEAD, encoding *base64.Encoding, value string) (string, error) {
	raw, err := encoding.DecodeString(value)
	if err != nil || len(raw) < aead.NonceSize() {
		return "", errors.New("invalid ciphertext")
	}
	plain, err := aead.Open(nil, raw[:aead.NonceSize()], raw[aead.NonceSize():], nil)
	return string(plain), err
}