package mongo

import (
	"context"
	"fmt"

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

	"bit.admedia.com/scm/ad/adpilot-indexing-code-parser.com/internal/service/parser"
)

const (
	collectionDocChunks           = "doc_chunks"
	collectionRepoDocumentSections = "repo_document_sections"
)

// ListDocChunks returns documentation chunks for a snapshot.
func (s *Store) ListDocChunks(ctx context.Context, repoID, snapshotID string) ([]parser.DocChunkRecord, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	filter := bson.M{
		"repo_id": repoID,
		"$or": []bson.M{
			{"snapshot_id": snapshotID},
			{"metadata.snapshot_id": snapshotID},
			{"metadata.snapshot.snapshot_id": snapshotID},
		},
	}

	cursor, err := s.docsCollection(collectionDocChunks).Find(ctx, filter)
	if err != nil {
		return nil, fmt.Errorf("list doc chunks: %w", err)
	}
	defer cursor.Close(ctx)

	records := make([]parser.DocChunkRecord, 0)
	for cursor.Next(ctx) {
		var doc struct {
			ChunkID       string `bson:"chunk_id"`
			DocPath       string `bson:"doc_path"`
			SectionTitle  string `bson:"section_title"`
			ChunkText     string `bson:"chunk_text"`
		}
		if err := cursor.Decode(&doc); err != nil {
			return nil, fmt.Errorf("decode doc chunk: %w", err)
		}
		records = append(records, parser.DocChunkRecord{
			ChunkID:      doc.ChunkID,
			DocPath:      doc.DocPath,
			SectionTitle: doc.SectionTitle,
			ChunkText:    doc.ChunkText,
		})
	}
	if err := cursor.Err(); err != nil {
		return nil, fmt.Errorf("iterate doc chunks: %w", err)
	}
	return records, nil
}

// ListDocSections returns documentation sections for docs in a repository.
func (s *Store) ListDocSections(ctx context.Context, repoID, snapshotID string) ([]parser.DocSectionRecord, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	// Sections are keyed by repo_id + doc_path; snapshot scope is implied by matching doc_paths from chunks.
	chunks, err := s.ListDocChunks(ctx, repoID, snapshotID)
	if err != nil {
		return nil, err
	}
	if len(chunks) == 0 {
		return nil, nil
	}

	docPaths := make(map[string]struct{})
	for _, chunk := range chunks {
		docPaths[chunk.DocPath] = struct{}{}
	}
	paths := make([]string, 0, len(docPaths))
	for path := range docPaths {
		paths = append(paths, path)
	}

	cursor, err := s.docsCollection(collectionRepoDocumentSections).Find(ctx, bson.M{
		"repo_id":  repoID,
		"doc_path": bson.M{"$in": paths},
	})
	if err != nil {
		return nil, fmt.Errorf("list doc sections: %w", err)
	}
	defer cursor.Close(ctx)

	records := make([]parser.DocSectionRecord, 0)
	for cursor.Next(ctx) {
		var doc struct {
			ID           string `bson:"_id"`
			DocPath      string `bson:"doc_path"`
			SectionTitle string `bson:"section_title"`
			Heading      string `bson:"heading"`
		}
		if err := cursor.Decode(&doc); err != nil {
			return nil, fmt.Errorf("decode doc section: %w", err)
		}
		title := doc.SectionTitle
		if title == "" {
			title = doc.Heading
		}
		records = append(records, parser.DocSectionRecord{
			SectionID:    doc.ID,
			DocPath:      doc.DocPath,
			SectionTitle: title,
		})
	}
	if err := cursor.Err(); err != nil {
		return nil, fmt.Errorf("iterate doc sections: %w", err)
	}
	return records, nil
}