package tools

import (
	"context"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	"golang.org/x/sync/semaphore"

	"git.warky.dev/wdevs/amcs/internal/ai"
	"git.warky.dev/wdevs/amcs/internal/session"
	"git.warky.dev/wdevs/amcs/internal/store"
)

const backfillConcurrency = 4

type BackfillTool struct {
	store      *store.DB
	embeddings *ai.EmbeddingRunner
	sessions   *session.ActiveProjects
	logger     *slog.Logger
}

type BackfillInput struct {
	Project         string `json:"project,omitempty" jsonschema:"optional project name or id to scope the backfill"`
	Limit           int    `json:"limit,omitempty" jsonschema:"maximum number of thoughts to process in one call; defaults to 100"`
	IncludeArchived bool   `json:"include_archived,omitempty" jsonschema:"whether to include archived thoughts; defaults to false"`
	OlderThanDays   int    `json:"older_than_days,omitempty" jsonschema:"only backfill thoughts older than N days; 0 means no restriction"`
	DryRun          bool   `json:"dry_run,omitempty" jsonschema:"report counts and sample ids without generating embeddings"`
}

type BackfillFailure struct {
	ID    string `json:"id"`
	Error string `json:"error"`
}

type BackfillOutput struct {
	Model    string            `json:"model"`
	Scanned  int               `json:"scanned"`
	Embedded int               `json:"embedded"`
	Skipped  int               `json:"skipped"`
	Failed   int               `json:"failed"`
	DryRun   bool              `json:"dry_run"`
	Failures []BackfillFailure `json:"failures,omitempty"`
}

func NewBackfillTool(db *store.DB, embeddings *ai.EmbeddingRunner, sessions *session.ActiveProjects, logger *slog.Logger) *BackfillTool {
	return &BackfillTool{store: db, embeddings: embeddings, sessions: sessions, logger: logger}
}

// QueueThought queues a single thought for background embedding generation.
// It is used by capture when the embedding provider is temporarily unavailable.
func (t *BackfillTool) QueueThought(ctx context.Context, id int64, content string) {
	go func() {
		started := time.Now()
		idStr := fmt.Sprint(id)
		t.logger.Info("background embedding started",
			slog.String("thought_id", idStr),
			slog.String("provider", t.embeddings.PrimaryProvider()),
			slog.String("model", t.embeddings.PrimaryModel()),
		)

		result, err := t.embeddings.Embed(ctx, content)
		if err != nil {
			t.logger.Warn("background embedding error",
				slog.String("thought_id", idStr),
				slog.String("provider", t.embeddings.PrimaryProvider()),
				slog.String("model", t.embeddings.PrimaryModel()),
				slog.String("stage", "embed"),
				slog.Duration("duration", time.Since(started)),
				slog.String("error", err.Error()),
			)
			return
		}
		if err := t.store.UpsertEmbedding(ctx, id, result.Model, result.Vector); err != nil {
			t.logger.Warn("background embedding error",
				slog.String("thought_id", idStr),
				slog.String("provider", t.embeddings.PrimaryProvider()),
				slog.String("model", result.Model),
				slog.String("stage", "upsert"),
				slog.Duration("duration", time.Since(started)),
				slog.String("error", err.Error()),
			)
			return
		}
		t.logger.Info("background embedding complete",
			slog.String("thought_id", idStr),
			slog.String("provider", t.embeddings.PrimaryProvider()),
			slog.String("model", result.Model),
			slog.Duration("duration", time.Since(started)),
		)
	}()
}

func (t *BackfillTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in BackfillInput) (*mcp.CallToolResult, BackfillOutput, error) {
	limit := in.Limit
	if limit <= 0 {
		limit = 100
	}

	project, err := resolveProject(ctx, t.store, t.sessions, req, in.Project, false)
	if err != nil {
		return nil, BackfillOutput{}, err
	}

	var projectID *int64
	if project != nil {
		projectID = &project.NumericID
	}

	primaryModel := t.embeddings.PrimaryModel()

	thoughts, err := t.store.ListThoughtsMissingEmbedding(ctx, primaryModel, limit, projectID, in.IncludeArchived, in.OlderThanDays)
	if err != nil {
		return nil, BackfillOutput{}, err
	}

	out := BackfillOutput{
		Model:   primaryModel,
		Scanned: len(thoughts),
		DryRun:  in.DryRun,
	}

	if in.DryRun || len(thoughts) == 0 {
		return nil, out, nil
	}

	start := time.Now()
	sem := semaphore.NewWeighted(backfillConcurrency)
	var mu sync.Mutex
	var wg sync.WaitGroup

	for _, thought := range thoughts {
		if ctx.Err() != nil {
			break
		}
		if err := sem.Acquire(ctx, 1); err != nil {
			break
		}
		wg.Add(1)
		go func(id int64, content string) {
			defer wg.Done()
			defer sem.Release(1)

			idStr := fmt.Sprint(id)
			result, embedErr := t.embeddings.Embed(ctx, content)
			if embedErr != nil {
				mu.Lock()
				out.Failures = append(out.Failures, BackfillFailure{ID: idStr, Error: embedErr.Error()})
				mu.Unlock()
				t.logger.Warn("backfill embed failed", slog.String("thought_id", idStr), slog.String("error", embedErr.Error()))
				return
			}

			if upsertErr := t.store.UpsertEmbedding(ctx, id, result.Model, result.Vector); upsertErr != nil {
				mu.Lock()
				out.Failures = append(out.Failures, BackfillFailure{ID: idStr, Error: upsertErr.Error()})
				mu.Unlock()
				t.logger.Warn("backfill upsert failed", slog.String("thought_id", idStr), slog.String("error", upsertErr.Error()))
				return
			}

			mu.Lock()
			out.Embedded++
			mu.Unlock()
		}(thought.ID, thought.Content)
	}

	wg.Wait()

	out.Failed = len(out.Failures)
	out.Skipped = out.Scanned - out.Embedded - out.Failed

	t.logger.Info("backfill completed",
		slog.String("model", primaryModel),
		slog.Int("scanned", out.Scanned),
		slog.Int("embedded", out.Embedded),
		slog.Int("failed", out.Failed),
		slog.Duration("duration", time.Since(start)),
	)

	return nil, out, nil
}
