package mongo

import (
	"context"
	"fmt"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

const (
	collectionCodeGraphArtifacts   = "code_graph_artifacts"
	collectionCodeChunks           = "code_chunks"
	collectionSnapshotGraphs       = "snapshot_graphs"
	collectionGraphDeltas          = "graph_deltas"
	collectionCommitProcessingState = "commit_processing_state"
)

// EnsureIndexes creates required MongoDB indexes.
func (s *Store) EnsureIndexes(ctx context.Context) error {
	if err := s.ensureCodeGraphArtifactIndexes(ctx); err != nil {
		return err
	}
	if err := s.ensureCodeGraphArtifactPartIndexes(ctx); err != nil {
		return err
	}
	if err := s.ensureCodeChunkIndexes(ctx); err != nil {
		return err
	}
	if err := s.ensureSnapshotGraphIndexes(ctx); err != nil {
		return err
	}
	if err := s.ensureSnapshotGraphPartIndexes(ctx); err != nil {
		return err
	}
	if err := s.ensureGraphDeltaIndexes(ctx); err != nil {
		return err
	}
	if err := s.ensureGraphDeltaPartIndexes(ctx); err != nil {
		return err
	}
	return s.ensureCommitProcessingStateIndexes(ctx)
}

func (s *Store) ensureCodeGraphArtifactIndexes(ctx context.Context) error {
	col := s.collection(collectionCodeGraphArtifacts)

	_, err := col.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys:    bson.D{{Key: "artifact_id", Value: 1}},
		Options: options.Index().SetUnique(true),
	})
	if err != nil {
		return fmt.Errorf("create code_graph_artifacts artifact_id index: %w", err)
	}

	_, err = col.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys: bson.D{
			{Key: "repo_id", Value: 1},
			{Key: "snapshot_id", Value: 1},
		},
	})
	if err != nil {
		return fmt.Errorf("create code_graph_artifacts repo snapshot index: %w", err)
	}

	_, err = col.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys: bson.D{
			{Key: "repo_id", Value: 1},
			{Key: "commit_sha", Value: 1},
			{Key: "file_path", Value: 1},
		},
	})
	if err != nil {
		return fmt.Errorf("create code_graph_artifacts commit file index: %w", err)
	}
	return nil
}

func (s *Store) ensureCodeGraphArtifactPartIndexes(ctx context.Context) error {
	col := s.collection(collectionCodeGraphArtifactParts)

	_, err := col.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys: bson.D{
			{Key: "artifact_id", Value: 1},
			{Key: "part_index", Value: 1},
		},
		Options: options.Index().SetUnique(true),
	})
	if err != nil {
		return fmt.Errorf("create code_graph_artifact_parts index: %w", err)
	}
	return nil
}

func (s *Store) ensureSnapshotGraphPartIndexes(ctx context.Context) error {
	col := s.collection(collectionSnapshotGraphParts)

	_, err := col.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys: bson.D{
			{Key: "artifact_id", Value: 1},
			{Key: "part_index", Value: 1},
		},
		Options: options.Index().SetUnique(true),
	})
	if err != nil {
		return fmt.Errorf("create snapshot_graph_parts index: %w", err)
	}
	return nil
}

func (s *Store) ensureCodeChunkIndexes(ctx context.Context) error {
	col := s.collection(collectionCodeChunks)

	_, err := col.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys: bson.D{
			{Key: "repo_id", Value: 1},
			{Key: "snapshot_id", Value: 1},
		},
	})
	if err != nil {
		return fmt.Errorf("create code_chunks repo snapshot index: %w", err)
	}

	_, err = col.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys: bson.D{
			{Key: "repo_id", Value: 1},
			{Key: "snapshot_id", Value: 1},
			{Key: "chunk_id", Value: 1},
		},
		Options: options.Index().SetUnique(true),
	})
	if err != nil {
		return fmt.Errorf("create code_chunks chunk_id index: %w", err)
	}
	return nil
}

func (s *Store) ensureSnapshotGraphIndexes(ctx context.Context) error {
	col := s.collection(collectionSnapshotGraphs)

	_, err := col.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys:    bson.D{{Key: "artifact_id", Value: 1}},
		Options: options.Index().SetUnique(true),
	})
	if err != nil {
		return fmt.Errorf("create snapshot_graphs artifact_id index: %w", err)
	}

	_, err = col.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys: bson.D{
			{Key: "repo_id", Value: 1},
			{Key: "snapshot_id", Value: 1},
		},
	})
	if err != nil {
		return fmt.Errorf("create snapshot_graphs repo snapshot index: %w", err)
	}
	return nil
}

func (s *Store) ensureGraphDeltaIndexes(ctx context.Context) error {
	col := s.collection(collectionGraphDeltas)

	_, err := col.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys:    bson.D{{Key: "delta_id", Value: 1}},
		Options: options.Index().SetUnique(true),
	})
	if err != nil {
		return fmt.Errorf("create graph_deltas delta_id index: %w", err)
	}

	_, err = col.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys: bson.D{
			{Key: "repo_id", Value: 1},
			{Key: "snapshot_id", Value: 1},
			{Key: "target_commit_sha", Value: 1},
		},
	})
	if err != nil {
		return fmt.Errorf("create graph_deltas snapshot commit index: %w", err)
	}
	return nil
}

func (s *Store) ensureGraphDeltaPartIndexes(ctx context.Context) error {
	col := s.collection(collectionGraphDeltaParts)

	_, err := col.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys: bson.D{
			{Key: "delta_id", Value: 1},
			{Key: "part_index", Value: 1},
		},
		Options: options.Index().SetUnique(true),
	})
	if err != nil {
		return fmt.Errorf("create graph_delta_parts index: %w", err)
	}
	return nil
}

func (s *Store) ensureCommitProcessingStateIndexes(ctx context.Context) error {
	col := s.collection(collectionCommitProcessingState)

	_, err := col.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys: bson.D{
			{Key: "repo_id", Value: 1},
			{Key: "snapshot_id", Value: 1},
			{Key: "commit_sha", Value: 1},
		},
		Options: options.Index().SetUnique(true),
	})
	if err != nil {
		return fmt.Errorf("create commit_processing_state commit index: %w", err)
	}
	return nil
}