package processor

import (
	"context"
	"screenmark/screenmark/.gen/haystack/haystack/model"
	"screenmark/screenmark/agents"
	"screenmark/screenmark/agents/client"
	"screenmark/screenmark/limits"
	"screenmark/screenmark/models"
	"sync"

	"github.com/charmbracelet/log"
)

const IMAGE_PROCESS_AT_A_TIME = 10

type ImageProcessor struct {
	imageModel models.ImageModel
	logger     *log.Logger

	descriptionAgent agents.DescriptionAgent
	stackAgent       client.AgentClient

	// TODO: add the notifier here

	Processor *Processor[model.Image]
}

func (p *ImageProcessor) setImageToProcess(ctx context.Context, image model.Image) {
	err := p.imageModel.UpdateProcess(ctx, image.ID, model.Progress_InProgress)
	if err != nil {
		// TODO: what can we actually do here for the errors?
		// We can't stop the work for the others

		p.logger.Error("failed to update image", "err", err)

		// TODO: we can use context here to actually pass some information through
		return
	}
}

func (p *ImageProcessor) describe(ctx context.Context, image model.Image) {
	descriptionSubLogger := p.logger.With("describe image", image.ID)

	err := p.descriptionAgent.Describe(descriptionSubLogger, image.ID, image.ImageName, image.Image)
	if err != nil {
		// Again, wtf do we do?
		// Although i think the agent actually returns an error when it's finished
		p.logger.Error("failed to describe image", "err", err)
		return
	}
}

func (p *ImageProcessor) extractInfo(ctx context.Context, image model.Image) {
	err := p.stackAgent.RunAgent(image.UserID, image.ID, image.ImageName, image.Image)
	if err != nil {
		// Again, wtf do we do?
		// Although i think the agent actually returns an error when it's finished
		p.logger.Error("failed to process image", "err", err)
		return
	}
}

func (p *ImageProcessor) processImage(image model.Image) {
	p.logger.Info("Processing image", "ID", image.ID)

	ctx := context.Background()

	p.setImageToProcess(ctx, image)

	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		p.describe(ctx, image)
		wg.Done()
	}()

	go func() {
		p.extractInfo(ctx, image)
		wg.Done()
	}()

	wg.Wait()
}

func NewImageProcessor(logger *log.Logger, imageModel models.ImageModel, listModel models.StackModel, limitsManager limits.LimitsManagerMethods) ImageProcessor {
	descriptionAgent := agents.NewDescriptionAgent(logger, imageModel)
	stackAgent := agents.NewListAgent(logger, listModel, limitsManager)

	imageProcessor := ImageProcessor{
		imageModel:       imageModel,
		logger:           logger,
		descriptionAgent: descriptionAgent,
		stackAgent:       stackAgent,
	}

	imageProcessor.Processor = NewProcessor(int(IMAGE_PROCESS_AT_A_TIME), imageProcessor.processImage)

	return imageProcessor
}