package colabauth

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"io"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

const maxAccounts = 16
const maxAccountBytes = 1 << 20

var errPoolFull = errors.New("account pool is full; remove an account first")

var accountIDPattern = regexp.MustCompile(`^acct_[A-Za-z0-9_-]{43}$`)

type accountStore struct {
	Version  int           `json:"version"`
	Accounts []credentials `json:"accounts"`
}

func (h *Handler) accountPath() string { return h.managedPath + ".accounts.json" }

func readPrivateJSON(path string, value any) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()
	info, err := f.Stat()
	if err != nil {
		return err
	}
	if !info.Mode().IsRegular() || info.Mode().Perm()&0o077 != 0 {
		return errors.New("account file is not private")
	}
	data, err := io.ReadAll(io.LimitReader(f, maxAccountBytes+1))
	if err != nil {
		return err
	}
	if len(data) > maxAccountBytes {
		return errors.New("account file too large")
	}
	return json.Unmarshal(data, value)
}

func (h *Handler) readAccounts() (accountStore, error) {
	var store accountStore
	err := readPrivateJSON(h.accountPath(), &store)
	if errors.Is(err, os.ErrNotExist) {
		return accountStore{Version: 1, Accounts: []credentials{}}, nil
	}
	if err != nil {
		return store, err
	}
	if store.Version != 1 || store.Accounts == nil || len(store.Accounts) > maxAccounts {
		return store, errors.New("invalid account store")
	}
	seen := map[string]bool{}
	for _, c := range store.Accounts {
		if !accountIDPattern.MatchString(c.AccountID) || seen[c.AccountID] || !validCredentials(c) {
			return store, errors.New("invalid account record")
		}
		seen[c.AccountID] = true
	}
	return store, nil
}

func validCredentials(c credentials) bool {
	return c.RefreshToken != "" && c.ClientID != "" && c.ClientSecret != ""
}

func newAccountID() (string, error) {
	var data [32]byte
	if _, err := rand.Read(data[:]); err != nil {
		return "", err
	}
	return "acct_" + base64.RawURLEncoding.EncodeToString(data[:]), nil
}

// import once, without changing the credential an older gateway would use.
func (h *Handler) initializeAccounts() error {
	h.storeMu.Lock()
	defer h.storeMu.Unlock()
	return h.initializeAccountsLocked()
}

func (h *Handler) initializeAccountsLocked() error {
	if _, err := os.Stat(h.accountPath()); err == nil {
		_, err = h.readAccounts()
		return err
	} else if !errors.Is(err, os.ErrNotExist) {
		return err
	}
	var old credentials
	if err := readPrivateJSON(h.managedPath, &old); errors.Is(err, os.ErrNotExist) {
		return nil
	} else if err != nil {
		return err
	}
	if !validCredentials(old) {
		return errors.New("invalid legacy credential")
	}
	id, err := newAccountID()
	if err != nil {
		return err
	}
	old.AccountID = id
	return writePrivateJSON(h.accountPath(), accountStore{Version: 1, Accounts: []credentials{old}})
}

func (h *Handler) save(creds credentials) error {
	if !validCredentials(creds) || creds.GoogleSubject == "" || creds.Email == "" {
		return errors.New("invalid account credential")
	}
	h.storeMu.Lock()
	defer h.storeMu.Unlock()
	if err := h.initializeAccountsLocked(); err != nil {
		return err
	}
	store, err := h.readAccounts()
	if err != nil {
		return err
	}
	for i, old := range store.Accounts {
		// older credentials have no Google subject. Upgrade that one verified email match.
		if old.GoogleSubject == creds.GoogleSubject || (old.GoogleSubject == "" && strings.EqualFold(old.Email, creds.Email)) {
			creds.AccountID = old.AccountID
			store.Accounts[i] = creds
			return writePrivateJSON(h.accountPath(), store)
		}
	}
	if len(store.Accounts) >= maxAccounts {
		return errPoolFull
	}
	creds.AccountID, err = newAccountID()
	if err != nil {
		return err
	}
	store.Accounts = append(store.Accounts, creds)
	return writePrivateJSON(h.accountPath(), store)
}

func (h *Handler) removeAccount(id string) (bool, error) {
	h.storeMu.Lock()
	defer h.storeMu.Unlock()
	store, err := h.readAccounts()
	if err != nil {
		return false, err
	}
	for i, account := range store.Accounts {
		if account.AccountID == id {
			store.Accounts = append(store.Accounts[:i], store.Accounts[i+1:]...)
			return true, writePrivateJSON(h.accountPath(), store)
		}
	}
	return false, nil
}

func writePrivateJSON(path string, value any) error {
	data, err := json.MarshalIndent(value, "", "  ")
	if err != nil {
		return err
	}
	if len(data)+1 > maxAccountBytes {
		return errors.New("account file too large")
	}
	dir := filepath.Dir(path)
	if err := ensurePrivateDir(dir); err != nil {
		return err
	}
	tmp, err := os.CreateTemp(dir, ".colab-credentials-*")
	if err != nil {
		return err
	}
	defer os.Remove(tmp.Name())
	defer tmp.Close()
	if err := tmp.Chmod(0o600); err != nil {
		return err
	}
	if _, err := tmp.Write(append(data, '\n')); err != nil {
		return err
	}
	if err := tmp.Sync(); err != nil {
		return err
	}
	if err := tmp.Close(); err != nil {
		return err
	}
	return os.Rename(tmp.Name(), path)
}