package mongo

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

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

	"bit.admedia.com/scm/ad/adpilot-indexing-code-parser.com/internal/contracts/graph"
	"bit.admedia.com/scm/ad/adpilot-indexing-code-parser.com/internal/store"
)

var _ store.ArtifactStore = (*Store)(nil)

type codeGraphArtifactDoc struct {
	ArtifactID    string       `bson:"artifact_id"`
	RepoID        string       `bson:"repo_id"`
	SnapshotID    string       `bson:"snapshot_id"`
	CommitSHA     string       `bson:"commit_sha"`
	FilePath      string       `bson:"file_path"`
	SchemaVersion string       `bson:"schema_version"`
	Chunked       bool         `bson:"chunked,omitempty"`
	PartCount     int          `bson:"part_count,omitempty"`
	NodeCount     int          `bson:"node_count"`
	EdgeCount     int          `bson:"edge_count"`
	Nodes         []graph.Node `bson:"nodes,omitempty"`
	Edges         []graph.Edge `bson:"edges,omitempty"`
	UpdatedAt     time.Time    `bson:"updated_at"`
}

// SaveArtifact upserts a per-file code graph artifact.
func (s *Store) SaveArtifact(ctx context.Context, artifact graph.Artifact) (string, string, error) {
	if err := ctx.Err(); err != nil {
		return "", "", err
	}
	if err := validateArtifactMetadata(artifact); err != nil {
		return "", "", err
	}

	filePath, err := filePathFromArtifact(artifact)
	if err != nil {
		return "", "", err
	}

	id := store.BuildArtifactID(artifact.SnapshotID, artifact.RepoID, artifact.CommitSHA, filePath)
	if artifact.SchemaVersion == "" {
		artifact.SchemaVersion = graph.SchemaVersionV1
	}

	parts := splitGraphPayload(artifact.Nodes, artifact.Edges, s.maxPartBytes)
	if len(parts) == 1 {
		doc := codeGraphArtifactDoc{
			ArtifactID:    id,
			RepoID:        artifact.RepoID,
			SnapshotID:    artifact.SnapshotID,
			CommitSHA:     artifact.CommitSHA,
			FilePath:      filePath,
			SchemaVersion: artifact.SchemaVersion,
			Chunked:       false,
			PartCount:     0,
			NodeCount:     len(artifact.Nodes),
			EdgeCount:     len(artifact.Edges),
			Nodes:         parts[0].Nodes,
			Edges:         parts[0].Edges,
			UpdatedAt:     time.Now().UTC(),
		}
		if err := s.upsertCodeGraphManifest(ctx, doc); err != nil {
			if mongoIsDocumentTooLarge(err) {
				return s.saveChunkedCodeGraphArtifact(ctx, id, artifact, filePath, artifact.Nodes, artifact.Edges)
			}
			return "", "", fmt.Errorf("upsert code graph artifact: %w", err)
		}
		if err := s.deleteGraphParts(ctx, collectionCodeGraphArtifactParts, id); err != nil {
			return "", "", err
		}
		return id, store.MongoURI(collectionCodeGraphArtifacts, id), nil
	}

	return s.saveChunkedCodeGraphArtifact(ctx, id, artifact, filePath, artifact.Nodes, artifact.Edges)
}

func (s *Store) saveChunkedCodeGraphArtifact(
	ctx context.Context,
	id string,
	artifact graph.Artifact,
	filePath string,
	nodes []graph.Node,
	edges []graph.Edge,
) (string, string, error) {
	parts := splitGraphPayload(nodes, edges, s.maxPartBytes)
	if len(parts) == 1 {
		parts = splitGraphPayload(nodes, edges, s.maxPartBytes/2)
	}
	if len(parts) == 1 {
		return "", "", fmt.Errorf("upsert code graph artifact: graph payload exceeds mongodb document limit")
	}

	manifest := codeGraphArtifactDoc{
		ArtifactID:    id,
		RepoID:        artifact.RepoID,
		SnapshotID:    artifact.SnapshotID,
		CommitSHA:     artifact.CommitSHA,
		FilePath:      filePath,
		SchemaVersion: artifact.SchemaVersion,
		Chunked:       true,
		PartCount:     len(parts),
		NodeCount:     len(nodes),
		EdgeCount:     len(edges),
		UpdatedAt:     time.Now().UTC(),
	}
	if err := s.upsertCodeGraphManifest(ctx, manifest); err != nil {
		return "", "", fmt.Errorf("upsert code graph artifact manifest: %w", err)
	}
	if err := s.saveGraphParts(ctx, collectionCodeGraphArtifactParts, id, parts); err != nil {
		return "", "", fmt.Errorf("save code graph artifact parts: %w", err)
	}
	return id, store.MongoURI(collectionCodeGraphArtifacts, id), nil
}

func (s *Store) upsertCodeGraphManifest(ctx context.Context, doc codeGraphArtifactDoc) error {
	_, err := s.collection(collectionCodeGraphArtifacts).UpdateOne(
		ctx,
		bson.M{"artifact_id": doc.ArtifactID},
		bson.M{"$set": doc},
		options.Update().SetUpsert(true),
	)
	return err
}

// GetArtifact loads a graph artifact by ID.
func (s *Store) GetArtifact(ctx context.Context, artifactID string) (graph.Artifact, error) {
	var artifact graph.Artifact
	if err := ctx.Err(); err != nil {
		return artifact, err
	}
	if err := validateArtifactID(artifactID); err != nil {
		return artifact, err
	}

	doc, err := s.findCodeGraphManifest(ctx, bson.M{"artifact_id": artifactID})
	if err != nil {
		if errors.Is(err, store.ErrArtifactNotFound) {
			return artifact, fmt.Errorf("%w: %s", store.ErrArtifactNotFound, artifactID)
		}
		return artifact, err
	}
	return s.artifactFromCodeGraphDoc(ctx, doc)
}

// GetArtifactMetadata loads artifact summary fields.
func (s *Store) GetArtifactMetadata(ctx context.Context, artifactID string) (store.ArtifactMetadata, error) {
	if err := ctx.Err(); err != nil {
		return store.ArtifactMetadata{}, err
	}
	if err := validateArtifactID(artifactID); err != nil {
		return store.ArtifactMetadata{}, err
	}

	doc, err := s.findCodeGraphManifest(ctx, bson.M{"artifact_id": artifactID})
	if err != nil {
		return store.ArtifactMetadata{}, err
	}

	filePath := doc.FilePath
	if filePath == "" {
		filePath, err = filePathFromNodes(doc.Nodes)
		if err != nil {
			return store.ArtifactMetadata{}, err
		}
	}

	return store.ArtifactMetadata{
		ArtifactID:    artifactID,
		SchemaVersion: doc.SchemaVersion,
		RepoID:        doc.RepoID,
		SnapshotID:    doc.SnapshotID,
		CommitSHA:     doc.CommitSHA,
		FilePath:      filePath,
		NodeCount:     doc.nodeCount(),
		EdgeCount:     doc.edgeCount(),
		ArtifactURI:   store.MongoURI(collectionCodeGraphArtifacts, artifactID),
	}, nil
}

// FindArtifactByCommitFile loads an artifact by repo, commit, and file path.
func (s *Store) FindArtifactByCommitFile(ctx context.Context, repoID, commitSHA, filePath string) (graph.Artifact, string, error) {
	var empty graph.Artifact
	if err := ctx.Err(); err != nil {
		return empty, "", err
	}

	doc, err := s.findCodeGraphManifest(ctx, bson.M{
		"repo_id":    repoID,
		"commit_sha": commitSHA,
		"file_path":  filePath,
	})
	if err != nil {
		if errors.Is(err, store.ErrArtifactNotFound) {
			return empty, "", fmt.Errorf("%w: repo=%s commit=%s file=%s", store.ErrArtifactNotFound, repoID, commitSHA, filePath)
		}
		return empty, "", err
	}

	artifact, err := s.artifactFromCodeGraphDoc(ctx, doc)
	if err != nil {
		return empty, "", err
	}
	return artifact, doc.ArtifactID, nil
}

// FindArtifactBySnapshotFile loads a parsed file artifact for a snapshot (HEAD parse).
func (s *Store) FindArtifactBySnapshotFile(ctx context.Context, repoID, snapshotID, filePath string) (graph.Artifact, string, error) {
	var empty graph.Artifact
	if err := ctx.Err(); err != nil {
		return empty, "", err
	}

	doc, err := s.findCodeGraphManifest(ctx, bson.M{
		"repo_id":     repoID,
		"snapshot_id": snapshotID,
		"file_path":   filePath,
	})
	if err != nil {
		if errors.Is(err, store.ErrArtifactNotFound) {
			return empty, "", fmt.Errorf("%w: repo=%s snapshot=%s file=%s", store.ErrArtifactNotFound, repoID, snapshotID, filePath)
		}
		return empty, "", err
	}

	artifact, err := s.artifactFromCodeGraphDoc(ctx, doc)
	if err != nil {
		return empty, "", err
	}
	return artifact, doc.ArtifactID, nil
}

// ListArtifactsBySnapshot returns all per-file artifacts for a snapshot.
func (s *Store) ListArtifactsBySnapshot(ctx context.Context, repoID, snapshotID string) ([]graph.Artifact, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	cursor, err := s.collection(collectionCodeGraphArtifacts).Find(ctx, bson.M{
		"repo_id":     repoID,
		"snapshot_id": snapshotID,
	})
	if err != nil {
		return nil, fmt.Errorf("list artifacts by snapshot: %w", err)
	}
	defer cursor.Close(ctx)

	artifacts := make([]graph.Artifact, 0)
	for cursor.Next(ctx) {
		var doc codeGraphArtifactDoc
		if err := cursor.Decode(&doc); err != nil {
			return nil, fmt.Errorf("decode code graph artifact: %w", err)
		}
		artifact, err := s.artifactFromCodeGraphDoc(ctx, doc)
		if err != nil {
			return nil, err
		}
		artifacts = append(artifacts, artifact)
	}
	if err := cursor.Err(); err != nil {
		return nil, fmt.Errorf("iterate code graph artifacts: %w", err)
	}
	return artifacts, nil
}

// CountCodeArtifacts returns the number of saved per-file code graph artifacts
// for a snapshot. This is the exact set buildMergedSnapshotGraph consumes, so it
// is the authoritative "how much of the graph is actually ready to merge" signal
// — unlike the fan-in path/skip counters, which can race ahead of artifact saves
// and trip a premature, permanently-sealed truncated snapshot graph.
func (s *Store) CountCodeArtifacts(ctx context.Context, repoID, snapshotID string) (int, error) {
	if err := ctx.Err(); err != nil {
		return 0, err
	}
	count, err := s.collection(collectionCodeGraphArtifacts).CountDocuments(ctx, bson.M{
		"repo_id":     repoID,
		"snapshot_id": snapshotID,
	})
	if err != nil {
		return 0, fmt.Errorf("count code graph artifacts: %w", err)
	}
	return int(count), nil
}

func (s *Store) findCodeGraphManifest(ctx context.Context, filter bson.M) (codeGraphArtifactDoc, error) {
	var doc codeGraphArtifactDoc
	err := s.collection(collectionCodeGraphArtifacts).FindOne(ctx, filter).Decode(&doc)
	if err != nil {
		if errors.Is(err, mongo.ErrNoDocuments) {
			return doc, store.ErrArtifactNotFound
		}
		return doc, fmt.Errorf("find code graph artifact: %w", err)
	}
	return doc, nil
}

func (s *Store) artifactFromCodeGraphDoc(ctx context.Context, doc codeGraphArtifactDoc) (graph.Artifact, error) {
	nodes, edges := doc.Nodes, doc.Edges
	if resolveChunked(doc.Chunked, doc.PartCount, len(doc.Nodes), len(doc.Edges)) {
		parts, err := s.loadGraphParts(ctx, collectionCodeGraphArtifactParts, doc.ArtifactID, doc.PartCount)
		if err != nil {
			return graph.Artifact{}, fmt.Errorf("load code graph artifact parts: %w", err)
		}
		nodes, edges = reassembleGraphPayload(parts)
	}

	return docToArtifact(doc, nodes, edges), nil
}

func docToArtifact(doc codeGraphArtifactDoc, nodes []graph.Node, edges []graph.Edge) graph.Artifact {
	return graph.Artifact{
		SchemaVersion: doc.SchemaVersion,
		RepoID:        doc.RepoID,
		CommitSHA:     doc.CommitSHA,
		SnapshotID:    doc.SnapshotID,
		Nodes:         nodes,
		Edges:         edges,
	}
}

func (d codeGraphArtifactDoc) nodeCount() int {
	if d.NodeCount > 0 {
		return d.NodeCount
	}
	return len(d.Nodes)
}

func (d codeGraphArtifactDoc) edgeCount() int {
	if d.EdgeCount > 0 {
		return d.EdgeCount
	}
	return len(d.Edges)
}

func validateArtifactMetadata(artifact graph.Artifact) error {
	if strings.TrimSpace(artifact.RepoID) == "" {
		return errors.New("artifact repo_id is required")
	}
	if strings.TrimSpace(artifact.SnapshotID) == "" {
		return errors.New("artifact snapshot_id is required")
	}
	if strings.TrimSpace(artifact.CommitSHA) == "" {
		return errors.New("artifact commit_sha is required")
	}
	return nil
}

func filePathFromArtifact(artifact graph.Artifact) (string, error) {
	return filePathFromNodes(artifact.Nodes)
}

func filePathFromNodes(nodes []graph.Node) (string, error) {
	for _, node := range nodes {
		if node.Kind == graph.NodeKindFile && strings.TrimSpace(node.Path) != "" {
			return node.Path, nil
		}
	}
	return "", errors.New("artifact file node with path is required")
}

func validateArtifactID(id string) error {
	if strings.TrimSpace(id) == "" {
		return store.ErrInvalidArtifactID
	}
	if strings.Contains(id, "..") || strings.ContainsAny(id, `/\`) {
		return store.ErrInvalidArtifactID
	}
	return nil
}