package storage

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/hashicorp/go-memdb"
	"github.com/pingcap/kvproto/pkg/keyspacepb"
	"github.com/pingcap/log"
	"go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
	clientv3 "go.etcd.io/etcd/client/v3"
	"go.uber.org/zap"
)

const keyspaceStateIndex = "state"

func keyspaceMetaTableSchema() *memdb.DBSchema {
	return &memdb.DBSchema{
		Tables: map[string]*memdb.TableSchema{
			keyspaceMetaTableName: {
				Name: keyspaceMetaTableName,
				Indexes: map[string]*memdb.IndexSchema{
					"id": { // Add this as a primary index
						Name:    "id",
						Unique:  true,
						Indexer: &memdb.StringFieldIndex{Field: "KeyspaceID"}, // Choose the primary unique identifier
					},
					string(IndexKeyspaceID): {
						Name:    string(IndexKeyspaceID),
						Unique:  true,
						Indexer: &memdb.StringFieldIndex{Field: "KeyspaceID"},
					},
					string(IndexKeyspaceName): {
						Name:    string(IndexKeyspaceName),
						Unique:  true,
						Indexer: &memdb.StringFieldIndex{Field: "KeyspaceName"},
					},
					string(IndexClusterID): {
						Name:         string(IndexClusterID),
						Unique:       false,
						AllowMissing: true,
						Indexer:      &memdb.StringFieldIndex{Field: "ClusterId"},
					},
					string(IndexTenantID): {
						Name:         string(IndexTenantID),
						Unique:       false,
						AllowMissing: true,
						Indexer:      &memdb.StringFieldIndex{Field: "TenantID"},
					},
					string(IndexProjectID): {
						Name:         string(IndexProjectID),
						Unique:       false,
						AllowMissing: true,
						Indexer:      &memdb.StringFieldIndex{Field: "ProjectID"},
					},
					string(IndexBillingPlan): {
						Name:         string(IndexBillingPlan),
						Unique:       false,
						AllowMissing: true,
						Indexer:      keyspaceBillingPlanIndexer{},
					},
					keyspaceStateIndex: {
						Name:    keyspaceStateIndex,
						Unique:  false,
						Indexer: &memdb.IntFieldIndex{Field: "State"},
					},
				},
			},
		},
	}
}

func (m *memDBStorage) fetchAllKeyspace(ctx context.Context) error {
	start := time.Now()
	startKey := m.keyspaceMetaPrefix
	endKey := clientv3.GetPrefixRangeEnd(startKey)
	totalLoaded := 0
	for {
		resp, err := m.etcdCli.Get(ctx, startKey, clientv3.WithRange(endKey), clientv3.WithLimit(etcdBatchSize))
		if err != nil {
			return err
		}
		if len(resp.Kvs) == 0 {
			break
		}
		txn := m.keyspaceMemDB.Txn(true)
		for _, kv := range resp.Kvs {
			meta, err := getKeyspaceMeta(kv.Value)
			if err != nil {
				return err
			}
			err = txn.Insert(keyspaceMetaTableName, meta)
			if err != nil {
				return err
			}
			totalLoaded += 1
		}
		txn.Commit()

		if resp.Header.Revision > m.initRevision {
			m.initRevision = resp.Header.Revision
		}
		// Set the next starting key to the last key retrieved
		startKey = string(resp.Kvs[len(resp.Kvs)-1].Key) + "\x00"
		// Stop when there are no more keys left
		if !resp.More {
			break
		}
	}
	log.Info("[storage] existing keyspace meta loaded",
		zap.Int64("init revision", m.initRevision),
		zap.Int("total meta loaded", totalLoaded),
		zap.Duration("duration", time.Since(start)),
	)
	return nil
}

func (m *memDBStorage) watchKeyspaceChanges(ctx context.Context) {
LOOP:
	log.Info("[storage] starting watch changes")
	watchCh := m.etcdCli.Watch(
		ctx,
		m.keyspaceMetaPrefix,
		clientv3.WithPrefix(),
		clientv3.WithRev(m.initRevision),
		clientv3.WithProgressNotify(),
	)
	for watchResp := range watchCh {
		if err := watchResp.Err(); err != nil {
			if errors.Is(err, rpctypes.ErrCompacted) {
				log.Fatal("[storage] watch changes failed", zap.Error(err))
			}
			log.Error("[storage] watch changes failed, reconnect after 10 seconds", zap.Error(err))
			time.Sleep(10 * time.Second) //nolint: gomnd // not need
			goto LOOP
		}
		if watchResp.IsProgressNotify() {
			m.initRevision = watchResp.Header.GetRevision()
			continue
		}
		txn := m.keyspaceMemDB.Txn(true)
		updated := 0
		for _, ev := range watchResp.Events {
			meta, err := getKeyspaceMeta(ev.Kv.Value)
			if err != nil {
				log.Error("[storage] unmarshal meta failed", zap.Error(err), zap.String("key", string(ev.Kv.Key)))
				continue
			}
			switch ev.Type {
			case clientv3.EventTypePut:
				if err = txn.Insert(keyspaceMetaTableName, meta); err != nil {
					log.Error("[storage] insert meta failed", zap.Error(err), zap.Any("meta", meta))
					continue
				}
			case clientv3.EventTypeDelete:
				if err = txn.Delete(keyspaceMetaTableName, meta); err != nil {
					log.Error("[storage] delete meta failed", zap.Error(err), zap.Any("meta", meta))
					continue
				}
			}
			updated += 1
		}
		txn.Commit()
		revision := watchResp.Header.GetRevision()
		m.initRevision = revision
		log.Info("[storage] db updated", zap.Int("count", updated), zap.Int64("revision", revision))
	}
}

func (m *memDBStorage) LookupKeyspace(index index, key string, enabledOnly bool) ([]*KeyspaceMetaData, error) {
	txn := m.keyspaceMemDB.Txn(false)
	it, err := txn.Get(keyspaceMetaTableName, string(index), key)
	if err != nil {
		return nil, err
	}
	result := make([]*KeyspaceMetaData, 0)
	for obj := it.Next(); obj != nil; obj = it.Next() {
		p := obj.(*KeyspaceMetaData)
		if enabledOnly && p.State != keyspacepb.KeyspaceState_ENABLED {
			continue
		}
		result = append(result, p)
	}
	return result, nil
}

func (m *memDBStorage) ListKeyspaces(filter KeyspaceListFilter) ([]*KeyspaceMetaData, error) {
	items, _, err := m.listKeyspacesWithPagination(filter, 0, 0)
	return items, err
}

// ListKeyspacesPaginated returns items for the requested page and total count. When limit <= 0, returns all items (ignore offset).
func (m *memDBStorage) ListKeyspacesPaginated(filter KeyspaceListFilter, limit, offset int) ([]*KeyspaceMetaData, int, error) {
	return m.listKeyspacesWithPagination(filter, limit, offset)
}

func (m *memDBStorage) listKeyspacesWithPagination(filter KeyspaceListFilter, limit, offset int) ([]*KeyspaceMetaData, int, error) {
	indexName := "id_prefix"
	indexArgs := []any{""}
	if filter.HasState {
		indexName = keyspaceStateIndex
		indexArgs = []any{filter.State}
	}

	txn := m.keyspaceMemDB.Txn(false)
	it, err := txn.Get(keyspaceMetaTableName, indexName, indexArgs...)
	if err != nil {
		return nil, 0, err
	}
	items := make([]*KeyspaceMetaData, 0)
	total := 0
	for obj := it.Next(); obj != nil; obj = it.Next() {
		p := obj.(*KeyspaceMetaData)
		if !matchKeyspaceFilter(p, filter) {
			continue
		}
		if total >= offset && (limit <= 0 || len(items) < limit) {
			items = append(items, p)
		}
		total++
	}
	return items, total, nil
}

func matchKeyspaceFilter(p *KeyspaceMetaData, filter KeyspaceListFilter) bool {
	if filter.HasState && p.State != filter.State {
		return false
	}
	return true
}

type keyspaceBillingPlanIndexer struct{}

func (keyspaceBillingPlanIndexer) FromArgs(args ...any) ([]byte, error) {
	if len(args) != 1 {
		return nil, fmt.Errorf("invalid args for billing plan indexer")
	}
	billingPlan, ok := args[0].(string)
	if !ok {
		return nil, fmt.Errorf("invalid args for billing plan indexer")
	}
	return []byte(billingPlan + "\x00"), nil
}

func (keyspaceBillingPlanIndexer) FromObject(raw any) (bool, []byte, error) {
	meta := raw.(*KeyspaceMetaData)
	if meta.IsEssential() {
		return true, []byte(essentialPlanName + "\x00"), nil
	}
	return false, nil, nil
}