package grpc

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/chroma-core/chroma/go/pkg/grpcutils"

	"github.com/chroma-core/chroma/go/pkg/common"
	"github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb"
	"github.com/chroma-core/chroma/go/pkg/sysdb/coordinator/model"
	"github.com/chroma-core/chroma/go/pkg/types"
	"github.com/pingcap/log"
	"go.uber.org/zap"
	"google.golang.org/protobuf/types/known/emptypb"
)

func (s *Server) ResetState(context.Context, *emptypb.Empty) (*coordinatorpb.ResetStateResponse, error) {
	log.Info("reset state")
	res := &coordinatorpb.ResetStateResponse{}
	err := s.coordinator.ResetState(context.Background())
	if err != nil {
		log.Error("error resetting state", zap.Error(err))
		return res, grpcutils.BuildInternalGrpcError(err.Error())
	}
	return res, nil
}

// Cases for get_or_create

// Case 0
// new_metadata is none, coll is an existing collection
// get_or_create should return the existing collection with existing metadata
// Essentially - an update with none is a no-op

// Case 1
// new_metadata is none, coll is a new collection
// get_or_create should create a new collection with the metadata of None

// Case 2
// new_metadata is not none, coll is an existing collection
// get_or_create should return the existing collection with updated metadata

// Case 3
// new_metadata is not none, coll is a new collection
// get_or_create should create a new collection with the new metadata, ignoring
// the metdata of in the input coll.

// The fact that we ignore the metadata of the generated collections is a
// bit weird, but it is the easiest way to excercise all cases

// NOTE: In current implementation, we do not support updating the metadata of an existing collection via this RPC.
//
// The call will fail if the collection already exists. Leaving the comments about cases 0,1,2,3 above for future reference.
// Refer to these issues for more context:
// https://github.com/chroma-core/chroma/issues/2390
// https://github.com/chroma-core/chroma/pull/2810
func (s *Server) CreateCollection(ctx context.Context, req *coordinatorpb.CreateCollectionRequest) (*coordinatorpb.CreateCollectionResponse, error) {
	res := &coordinatorpb.CreateCollectionResponse{}

	log.Info("CreateCollectionRequest", zap.Any("request", req))

	createCollection, err := convertToCreateCollectionModel(req)
	if err != nil {
		log.Error("CreateCollection failed. error converting to create collection model", zap.Error(err), zap.String("collection_id", req.Id), zap.String("collection_name", req.Name))
		res.Collection = &coordinatorpb.Collection{
			Id:                   req.Id,
			Name:                 req.Name,
			ConfigurationJsonStr: req.ConfigurationJsonStr,
			Dimension:            req.Dimension,
			Metadata:             req.Metadata,
			Tenant:               req.Tenant,
			Database:             req.Database,
		}
		res.Created = false
		return res, grpcutils.BuildInternalGrpcError(err.Error())
	}

	// Convert the request segments to create segment models
	createSegments := []*model.CreateSegment{}
	for _, segment := range req.Segments {
		createSegment, err := convertSegmentToModel(segment)
		if err != nil {
			log.Error("Error in creating segments for the collection", zap.Error(err))
			res.Collection = nil // We don't need to set the collection in case of error
			res.Created = false
			if err == common.ErrSegmentUniqueConstraintViolation {
				log.Error("segment id already exist", zap.Error(err))
				return res, grpcutils.BuildAlreadyExistsGrpcError(err.Error())
			}
			return res, grpcutils.BuildInternalGrpcError(err.Error())
		}
		createSegments = append(createSegments, createSegment)
	}

	// Create the collection and segments
	collection, created, err := s.coordinator.CreateCollectionAndSegments(ctx, createCollection, createSegments)
	if err != nil {
		log.Error("CreateCollection failed. error creating collection", zap.Error(err), zap.String("collection_id", req.Id), zap.String("collection_name", req.Name))
		res.Collection = &coordinatorpb.Collection{
			Id:                   req.Id,
			Name:                 req.Name,
			ConfigurationJsonStr: req.ConfigurationJsonStr,
			Dimension:            req.Dimension,
			Metadata:             req.Metadata,
			Tenant:               req.Tenant,
			Database:             req.Database,
		}
		res.Created = false
		if err == common.ErrCollectionUniqueConstraintViolation {
			return res, grpcutils.BuildAlreadyExistsGrpcError(err.Error())
		}
		return res, grpcutils.BuildInternalGrpcError(err.Error())
	}
	res.Collection = convertCollectionToProto(collection)
	res.Created = created
	log.Info("CreateCollection finished.", zap.String("collection_id", req.Id), zap.String("collection_name", req.Name), zap.Bool("created", created))
	return res, nil
}

func (s *Server) GetCollections(ctx context.Context, req *coordinatorpb.GetCollectionsRequest) (*coordinatorpb.GetCollectionsResponse, error) {
	collectionID := req.Id
	collectionName := req.Name
	tenantID := req.Tenant
	databaseName := req.Database
	limit := req.Limit
	offset := req.Offset

	res := &coordinatorpb.GetCollectionsResponse{}

	parsedCollectionID, err := types.ToUniqueID(collectionID)
	if err != nil {
		log.Error("GetCollections failed. collection id format error", zap.Error(err), zap.Stringp("collection_id", collectionID), zap.Stringp("collection_name", collectionName))
		return res, grpcutils.BuildInternalGrpcError(err.Error())
	}

	collections, err := s.coordinator.GetCollections(ctx, parsedCollectionID, collectionName, tenantID, databaseName, limit, offset)
	if err != nil {
		log.Error("GetCollections failed. ", zap.Error(err), zap.Stringp("collection_id", collectionID), zap.Stringp("collection_name", collectionName))
		return res, grpcutils.BuildInternalGrpcError(err.Error())
	}
	res.Collections = make([]*coordinatorpb.Collection, 0, len(collections))
	for _, collection := range collections {
		collectionpb := convertCollectionToProto(collection)
		res.Collections = append(res.Collections, collectionpb)
	}
	return res, nil
}

func (s *Server) CheckCollections(ctx context.Context, req *coordinatorpb.CheckCollectionsRequest) (*coordinatorpb.CheckCollectionsResponse, error) {
	res := &coordinatorpb.CheckCollectionsResponse{}
	res.Deleted = make([]bool, len(req.CollectionIds))

	for i, collectionID := range req.CollectionIds {
		parsedId, err := types.ToUniqueID(&collectionID)
		if err != nil {
			log.Error("CheckCollection failed. collection id format error", zap.Error(err), zap.String("collection_id", collectionID))
			return nil, grpcutils.BuildInternalGrpcError(err.Error())
		}
		deleted, err := s.coordinator.CheckCollection(ctx, parsedId)

		if err != nil {
			log.Error("CheckCollection failed", zap.Error(err), zap.String("collection_id", collectionID))
			return nil, grpcutils.BuildInternalGrpcError(err.Error())
		}

		res.Deleted[i] = deleted
	}
	return res, nil
}

func (s *Server) GetCollectionWithSegments(ctx context.Context, req *coordinatorpb.GetCollectionWithSegmentsRequest) (*coordinatorpb.GetCollectionWithSegmentsResponse, error) {
	collectionID := req.Id

	res := &coordinatorpb.GetCollectionWithSegmentsResponse{}

	parsedCollectionID, err := types.ToUniqueID(&collectionID)
	if err != nil {
		log.Error("GetCollectionWithSegments failed. collection id format error", zap.Error(err), zap.String("collection_id", collectionID))
		return res, grpcutils.BuildInternalGrpcError(err.Error())
	}

	collection, segments, err := s.coordinator.GetCollectionWithSegments(ctx, parsedCollectionID)
	if err != nil {
		log.Error("GetCollectionWithSegments failed. ", zap.Error(err), zap.String("collection_id", collectionID))
		return res, grpcutils.BuildInternalGrpcError(err.Error())
	}

	res.Collection = convertCollectionToProto(collection)
	segmentpbList := make([]*coordinatorpb.Segment, 0, len(segments))
	scopeToSegmentMap := map[coordinatorpb.SegmentScope]*coordinatorpb.Segment{}
	for _, segment := range segments {
		segmentpb := convertSegmentToProto(segment)
		scopeToSegmentMap[segmentpb.GetScope()] = segmentpb
		segmentpbList = append(segmentpbList, segmentpb)
	}

	if len(segmentpbList) != 3 {
		log.Error("GetCollectionWithSegments failed. Unexpected number of collection segments", zap.String("collection_id", collectionID))
		return res, grpcutils.BuildInternalGrpcError(fmt.Sprintf("Unexpected number of segments for collection %s: %d", collectionID, len(segmentpbList)))
	}

	scopes := []coordinatorpb.SegmentScope{coordinatorpb.SegmentScope_METADATA, coordinatorpb.SegmentScope_RECORD, coordinatorpb.SegmentScope_VECTOR}

	for _, scope := range scopes {
		if _, exists := scopeToSegmentMap[scope]; !exists {
			log.Error("GetCollectionWithSegments failed. Collection segment scope not found", zap.String("collection_id", collectionID), zap.String("missing_scope", scope.String()))
			return res, grpcutils.BuildInternalGrpcError(fmt.Sprintf("Missing segment scope for collection %s: %s", collectionID, scope.String()))
		}
	}

	res.Segments = segmentpbList
	return res, nil
}

func (s *Server) DeleteCollection(ctx context.Context, req *coordinatorpb.DeleteCollectionRequest) (*coordinatorpb.DeleteCollectionResponse, error) {
	collectionID := req.GetId()
	res := &coordinatorpb.DeleteCollectionResponse{}
	parsedCollectionID, err := types.Parse(collectionID)
	if err != nil {
		log.Error("DeleteCollection failed", zap.Error(err), zap.String("collection_id", collectionID))
		return res, grpcutils.BuildInternalGrpcError(err.Error())
	}
	deleteCollection := &model.DeleteCollection{
		ID:           parsedCollectionID,
		TenantID:     req.GetTenant(),
		DatabaseName: req.GetDatabase(),
	}
	err = s.coordinator.DeleteCollection(ctx, deleteCollection)
	if err != nil {
		log.Error("DeleteCollection failed", zap.Error(err), zap.String("collection_id", collectionID))
		if err == common.ErrCollectionDeleteNonExistingCollection {
			return res, grpcutils.BuildNotFoundGrpcError(err.Error())
		}
		return res, grpcutils.BuildInternalGrpcError(err.Error())
	}
	log.Info("DeleteCollection succeeded", zap.String("collection_id", collectionID))
	return res, nil
}

func (s *Server) UpdateCollection(ctx context.Context, req *coordinatorpb.UpdateCollectionRequest) (*coordinatorpb.UpdateCollectionResponse, error) {
	res := &coordinatorpb.UpdateCollectionResponse{}

	collectionID := req.Id
	parsedCollectionID, err := types.ToUniqueID(&collectionID)
	if err != nil {
		log.Error("UpdateCollection failed. collection id format error", zap.Error(err), zap.String("collection_id", collectionID))
		return res, grpcutils.BuildInternalGrpcError(err.Error())
	}

	updateCollection := &model.UpdateCollection{
		ID:        parsedCollectionID,
		Name:      req.Name,
		Dimension: req.Dimension,
	}

	resetMetadata := req.GetResetMetadata()
	updateCollection.ResetMetadata = resetMetadata
	metadata := req.GetMetadata()
	// Case 1: if resetMetadata is true, then delete all metadata for the collection
	// Case 2: if resetMetadata is true and metadata is not nil -> THIS SHOULD NEVER HAPPEN
	// Case 3: if resetMetadata is false, and the metadata is not nil - set the metadata to the value in metadata
	// Case 4: if resetMetadata is false and metadata is nil, then leave the metadata as is
	if resetMetadata {
		if metadata != nil {
			log.Error("UpdateCollection failed. reset metadata is true and metadata is not nil", zap.Any("metadata", metadata), zap.String("collection_id", collectionID))
			return res, grpcutils.BuildInternalGrpcError(common.ErrInvalidMetadataUpdate.Error())
		} else {
			updateCollection.Metadata = nil
		}
	} else {
		if metadata != nil {
			modelMetadata, err := convertCollectionMetadataToModel(metadata)
			if err != nil {
				log.Error("UpdateCollection failed. error converting collection metadata to model", zap.Error(err), zap.String("collection_id", collectionID))
				return res, grpcutils.BuildInternalGrpcError(err.Error())
			}
			updateCollection.Metadata = modelMetadata
		} else {
			updateCollection.Metadata = nil
		}
	}

	_, err = s.coordinator.UpdateCollection(ctx, updateCollection)

	if err != nil {
		log.Error("UpdateCollection failed. error updating collection", zap.Error(err), zap.String("collection_id", collectionID))
		if err == common.ErrCollectionUniqueConstraintViolation {
			return res, grpcutils.BuildAlreadyExistsGrpcError(err.Error())
		}
		return res, grpcutils.BuildInternalGrpcError(err.Error())
	}

	return res, nil
}

func (s *Server) FlushCollectionCompaction(ctx context.Context, req *coordinatorpb.FlushCollectionCompactionRequest) (*coordinatorpb.FlushCollectionCompactionResponse, error) {
	_, err := json.Marshal(req)
	if err != nil {
		log.Error("FlushCollectionCompaction failed. error marshalling request", zap.Error(err), zap.String("collection_id", req.CollectionId), zap.Int32("collection_version", req.CollectionVersion), zap.Int64("log_position", req.LogPosition))
		return nil, grpcutils.BuildInternalGrpcError(err.Error())
	}
	collectionID, err := types.ToUniqueID(&req.CollectionId)
	err = grpcutils.BuildErrorForUUID(collectionID, "collection", err)
	if err != nil {
		log.Error("FlushCollectionCompaction failed. error parsing collection id", zap.Error(err), zap.String("collection_id", req.CollectionId), zap.Int32("collection_version", req.CollectionVersion), zap.Int64("log_position", req.LogPosition))
		return nil, grpcutils.BuildInternalGrpcError(err.Error())
	}
	segmentCompactionInfo := make([]*model.FlushSegmentCompaction, 0, len(req.SegmentCompactionInfo))
	for _, flushSegmentCompaction := range req.SegmentCompactionInfo {
		segmentID, err := types.ToUniqueID(&flushSegmentCompaction.SegmentId)
		err = grpcutils.BuildErrorForUUID(segmentID, "segment", err)
		if err != nil {
			log.Error("FlushCollectionCompaction failed. error parsing segment id", zap.Error(err), zap.String("collection_id", req.CollectionId), zap.Int32("collection_version", req.CollectionVersion), zap.Int64("log_position", req.LogPosition))
			return nil, grpcutils.BuildInternalGrpcError(err.Error())
		}
		filePaths := make(map[string][]string)
		for key, filePath := range flushSegmentCompaction.FilePaths {
			filePaths[key] = filePath.Paths
		}
		segmentCompactionInfo = append(segmentCompactionInfo, &model.FlushSegmentCompaction{
			ID:        segmentID,
			FilePaths: filePaths,
		})
	}
	FlushCollectionCompaction := &model.FlushCollectionCompaction{
		ID:                       collectionID,
		TenantID:                 req.TenantId,
		LogPosition:              req.LogPosition,
		CurrentCollectionVersion: req.CollectionVersion,
		FlushSegmentCompactions:  segmentCompactionInfo,
	}
	flushCollectionInfo, err := s.coordinator.FlushCollectionCompaction(ctx, FlushCollectionCompaction)
	if err != nil {
		log.Error("FlushCollectionCompaction failed", zap.Error(err), zap.String("collection_id", req.CollectionId), zap.Int32("collection_version", req.CollectionVersion), zap.Int64("log_position", req.LogPosition))
		return nil, grpcutils.BuildInternalGrpcError(err.Error())
	}
	res := &coordinatorpb.FlushCollectionCompactionResponse{
		CollectionId:       flushCollectionInfo.ID,
		CollectionVersion:  flushCollectionInfo.CollectionVersion,
		LastCompactionTime: flushCollectionInfo.TenantLastCompactionTime,
	}
	return res, nil
}
