package main

import (
	"context"
	"flag"
	"math"
	"strings"
	"time"

	"cloud.google.com/go/firestore"
	"go.skia.org/infra/golden/go/fs_utils"
	"google.golang.org/api/iterator"

	ifirestore "go.skia.org/infra/go/firestore"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/golden/go/expectations"
	"go.skia.org/infra/golden/go/types"
)

const (
	maxOperationTime = 2 * time.Minute
	maxRetries       = 5

	dryrun = true
)

func main() {
	var (
		fsProjectID    = flag.String("fs_project_id", "skia-firestore", "The project with the firestore instance. Datastore and Firestore can't be in the same project.")
		oldFSNamespace = flag.String("old_fs_namespace", "", "Typically the instance id. e.g. 'chrome-gpu', 'skia', etc")
		newFSNamespace = flag.String("new_fs_namespace", "", "Typically the instance id. e.g. 'chrome'")
	)
	flag.Parse()

	if *oldFSNamespace == "" {
		sklog.Fatalf("You must include fs_namespace")
	}

	if *newFSNamespace == "" {
		sklog.Fatalf("You must include sql_db_name")
	}

	fsClient, err := ifirestore.NewClient(context.Background(), *fsProjectID, "gold", *oldFSNamespace, nil)
	if err != nil {
		sklog.Fatalf("Unable to configure Firestore: %s", err)
	}
	ctx := context.Background()
	oldNamespace := v3Impl{client: fsClient}

	fsClient, err = ifirestore.NewClient(context.Background(), *fsProjectID, "gold", *newFSNamespace, nil)
	if err != nil {
		sklog.Fatalf("Unable to configure Firestore: %s", err)
	}
	newNamespace := v3Impl{client: fsClient}

	// Fetch triage records
	records, err := oldNamespace.fetchTriageRecords(ctx)
	if err != nil {
		sklog.Fatalf("Fetching triage records: %s", err)
	}
	sklog.Infof("Should migrate %d records", len(records))

	if err := newNamespace.storeRecords(ctx, records); err != nil {
		sklog.Fatalf("Storing triage records: %s", err)
	}

	// Fetch Deltas
	deltas, err := oldNamespace.fetchExpectationDeltas(ctx)
	if err != nil {
		sklog.Fatalf("Fetching expectation deltas: %s", err)
	}

	if err := newNamespace.storeExpectationChanges(ctx, deltas); err != nil {
		sklog.Fatalf("Storing triage records: %s", err)
	}

	// Fetch entries
	exp, err := oldNamespace.fetchExpectations(ctx)
	if err != nil {
		sklog.Fatalf("Fetching expectation entries: %s", err)
	}
	if err := newNamespace.storeEntries(ctx, exp); err != nil {
		sklog.Fatalf("Storing triage records: %s", err)
	}
	sklog.Infof("Done")
}

const (
	v3Partitions         = "expstore_partitions_v3"
	v3ExpectationEntries = "entries"
	v3RecordEntries      = "triage_records"
	v3ChangeEntries      = "triage_changes"

	v3MasterPartition = "master"
	v3BeginningOfTime = 0
	v3EndOfTime       = math.MaxInt32

	v3DigestField = "digest"
)

type v3Impl struct {
	client *ifirestore.Client
}

type v3ExpectationEntry struct {
	Grouping types.TestName  `firestore:"grouping"`
	Digest   types.Digest    `firestore:"digest"`
	Updated  time.Time       `firestore:"updated"`
	LastUsed time.Time       `firestore:"last_used"`
	Ranges   []v3TriageRange `firestore:"ranges"`
	NeedsGC  bool            `firestore:"needs_gc"`
}

func (e *v3ExpectationEntry) id() string {
	s := string(e.Grouping) + "|" + string(e.Digest)
	// firestore gets cranky if there are / in key names
	return strings.Replace(s, "/", "-", -1)
}

type v3TriageRange struct {
	FirstIndex int                   `firestore:"first_index"`
	LastIndex  int                   `firestore:"last_index"`
	Label      expectations.LabelInt `firestore:"label"`
}

type v3TriageRecord struct {
	ID        string
	UserName  string    `firestore:"user"`
	TS        time.Time `firestore:"ts"`
	Changes   int       `firestore:"changes"`
	Committed bool      `firestore:"committed"`
}

type v3ExpectationChange struct {
	// RecordID refers to a document in the records collection.
	RecordID      string                `firestore:"record_id"`
	Grouping      types.TestName        `firestore:"grouping"`
	Digest        types.Digest          `firestore:"digest"`
	AffectedRange v3TriageRange         `firestore:"affected_range"`
	LabelBefore   expectations.LabelInt `firestore:"label_before"`
}

func (v v3Impl) fetchTriageRecords(ctx context.Context) (map[string][]v3TriageRecord, error) {
	// maps partition to records
	rv := map[string][]v3TriageRecord{}
	partitionIterator := v.client.Collection(v3Partitions).DocumentRefs(ctx)
	p, err := partitionIterator.Next()
	for ; err == nil; p, err = partitionIterator.Next() {
		var records []v3TriageRecord
		partition := p.ID
		sklog.Infof("Partition %s", partition)
		recordIterator := v.client.Collection(v3Partitions).Doc(partition).Collection(v3RecordEntries).Documents(ctx)
		docs, err := recordIterator.GetAll()
		if err != nil {
			return nil, skerr.Wrapf(err, "getting records for %s", partition)
		}
		for _, doc := range docs {
			var r v3TriageRecord
			if err := doc.DataTo(&r); err != nil {
				if err != nil {
					sklog.Warning("Corrupt triage record with id %s", doc.Ref.ID)
					continue
				}
			}
			r.ID = doc.Ref.ID
			records = append(records, r)
		}
		rv[partition] = records
	}
	if err != iterator.Done {
		return nil, skerr.Wrap(err)
	}

	return rv, nil
}

func (v v3Impl) storeRecords(ctx context.Context, recordsByPartition map[string][]v3TriageRecord) error {
	if dryrun {
		sklog.Infof("Would store %d partition of records", len(recordsByPartition))
		return nil
	}
	for partition, records := range recordsByPartition {
		sklog.Infof("Writing %d records for partition %s", len(records), partition)
		entryCollection := v.client.Collection(v3Partitions).Doc(partition).Collection(v3RecordEntries)

		err := v.client.BatchWrite(ctx, len(records), ifirestore.MAX_TRANSACTION_DOCS, maxOperationTime, nil, func(b *firestore.WriteBatch, i int) error {
			record := records[i]
			doc := entryCollection.Doc(record.ID)
			b.Set(doc, record)
			return nil
		})
		if err != nil {
			return skerr.Wrapf(err, "storing to partition %s", partition)
		}
	}
	return nil
}

func (v v3Impl) fetchExpectationDeltas(ctx context.Context) (map[string][]v3ExpectationChange, error) {
	rv := map[string][]v3ExpectationChange{} // Maps partition -> entries

	partitionIterator := v.client.Collection(v3Partitions).DocumentRefs(ctx)
	p, err := partitionIterator.Next()
	for ; err == nil; p, err = partitionIterator.Next() {
		partition := p.ID
		const numShards = 16
		base := v.client.Collection(v3Partitions).Doc(partition).Collection(v3ChangeEntries)

		queries := fs_utils.ShardOnDigest(base, v3DigestField, numShards)
		shardedEntries := make([][]v3ExpectationChange, numShards)

		err := v.client.IterDocsInParallel(ctx, "loadExpectationDeltas", partition, queries, maxRetries, maxOperationTime, func(i int, doc *firestore.DocumentSnapshot) error {
			if doc == nil {
				return nil
			}
			entry := v3ExpectationChange{}
			if err := doc.DataTo(&entry); err != nil {
				id := doc.Ref.ID
				return skerr.Wrapf(err, "corrupt data in firestore, could not unmarshal expectationChange with id %s", id)
			}

			shardedEntries[i] = append(shardedEntries[i], entry)
			return nil
		})

		if err != nil {
			return nil, skerr.Wrapf(err, "fetching expectation deltas for partition %s", partition)
		}

		var combinedEntries []v3ExpectationChange
		for _, shard := range shardedEntries {
			combinedEntries = append(combinedEntries, shard...)
		}

		rv[partition] = combinedEntries
	}
	return rv, nil
}

func (v v3Impl) storeExpectationChanges(ctx context.Context, changesByPartition map[string][]v3ExpectationChange) error {
	if dryrun {
		sklog.Infof("Would store %d partition of changes", len(changesByPartition))
		return nil
	}
	for partition, changes := range changesByPartition {
		sklog.Infof("Writing %d changes for partition %s", len(changes), partition)
		changesCollection := v.client.Collection(v3Partitions).Doc(partition).Collection(v3ChangeEntries)

		err := v.client.BatchWrite(ctx, len(changes), ifirestore.MAX_TRANSACTION_DOCS, maxOperationTime, nil, func(b *firestore.WriteBatch, i int) error {
			change := changes[i]
			doc := changesCollection.NewDoc()
			b.Set(doc, change)
			return nil
		})
		if err != nil {
			return skerr.Wrapf(err, "storing to partition %s", partition)
		}
	}
	return nil
}

func (v v3Impl) fetchExpectations(ctx context.Context) (map[string][]v3ExpectationEntry, error) {
	rv := map[string][]v3ExpectationEntry{} // Maps partition -> entries

	partitionIterator := v.client.Collection(v3Partitions).DocumentRefs(ctx)
	p, err := partitionIterator.Next()
	for ; err == nil; p, err = partitionIterator.Next() {
		partition := p.ID

		const numShards = 16
		base := v.client.Collection(v3Partitions).Doc(partition).Collection(v3ExpectationEntries)
		queries := fs_utils.ShardOnDigest(base, v3DigestField, numShards)
		shardedEntries := make([][]v3ExpectationEntry, numShards)

		err := v.client.IterDocsInParallel(ctx, "loadExpectations", partition, queries, maxRetries, maxOperationTime, func(i int, doc *firestore.DocumentSnapshot) error {
			if doc == nil {
				return nil
			}
			entry := v3ExpectationEntry{}
			if err := doc.DataTo(&entry); err != nil {
				id := doc.Ref.ID
				return skerr.Wrapf(err, "corrupt data in firestore, could not unmarshal expectationEntry with id %s", id)
			}
			if len(entry.Ranges) == 0 {
				// This should never happen, but we'll ignore these malformed entries if they do.
				return nil
			}
			shardedEntries[i] = append(shardedEntries[i], entry)
			return nil
		})

		if err != nil {
			return nil, skerr.Wrapf(err, "fetching expectations for partition %s", partition)
		}

		var combinedEntries []v3ExpectationEntry
		for _, shard := range shardedEntries {
			combinedEntries = append(combinedEntries, shard...)
		}

		rv[partition] = combinedEntries
		sklog.Infof("Fetched %d entries for partition %s", len(combinedEntries), partition)
	}
	return rv, nil
}

func (v v3Impl) storeEntries(ctx context.Context, entriesByPartition map[string][]v3ExpectationEntry) error {
	if dryrun {
		sklog.Infof("Would store %d partition of entries", len(entriesByPartition))
		return nil
	}
	for partition, entries := range entriesByPartition {
		sklog.Infof("Writing %d entries for partition %s", len(entries), partition)
		entryCollection := v.client.Collection(v3Partitions).Doc(partition).Collection(v3ExpectationEntries)

		err := v.client.BatchWrite(ctx, len(entries), ifirestore.MAX_TRANSACTION_DOCS, maxOperationTime, nil, func(b *firestore.WriteBatch, i int) error {
			entry := entries[i]
			doc := entryCollection.Doc(entry.id())
			b.Set(doc, entry)
			return nil
		})
		if err != nil {
			return skerr.Wrapf(err, "storing to partition %s", partition)
		}
	}
	return nil
}