package memoir

import (
	"context"
	"data-terminal-cloud/internal/dify"
	"data-terminal-cloud/internal/extraction"
	"sync"
)

type BookWorkflow interface {
	Generate(context.Context, []Source, func(Document) error) (Document, [][]byte, error)
}
type DifyWorkflow struct {
	Client *dify.Client
	Bridge *dify.Bridge
}

func NewDifyWorkflow(base, key, bridgeURL, bridgeSecret string) (*DifyWorkflow, error) {
	client, err := dify.NewClient(base, key)
	if err != nil {
		return nil, err
	}
	return &DifyWorkflow{client, dify.NewBridge(bridgeURL, bridgeSecret)}, nil
}
func (g *DifyWorkflow) Generate(ctx context.Context, sources []Source, publish func(Document) error) (Document, [][]byte, error) {
	ticket, err := g.Bridge.Register(ctx, "memoir", map[string]any{"stories": sources})
	if err != nil {
		return Document{}, nil, err
	}
	defer g.Bridge.Forget(ticket)
	var response struct {
		Document  Document `json:"document"`
		ErrorCode string   `json:"error_code"`
	}
	published := false
	run, runErr := g.Client.Run(ctx, map[string]any{"context_id": ticket.ID, "step_token": ticket.Token}, func(e dify.Event) error {
		if e.Data.NodeID != "validate-book" || published {
			return nil
		}
		if err := g.Bridge.Result(ctx, ticket, &response); err != nil {
			return err
		}
		if err := ValidateDocument(response.Document, sources); err != nil {
			return err
		}
		response.Document.WorkflowRunID = e.WorkflowRunID
		if err := publish(response.Document); err != nil {
			return err
		}
		published = true
		return nil
	})
	if runErr != nil {
		_ = g.Bridge.Result(ctx, ticket, &response)
		if response.ErrorCode == "MODEL_ACCOUNT_OVERDUE" {
			return Document{}, nil, extraction.ErrModelAccountOverdue
		}
		return Document{}, nil, runErr
	}
	if err = g.Bridge.Result(ctx, ticket, &response); err != nil {
		return Document{}, nil, err
	}
	if err = ValidateDocument(response.Document, sources); err != nil {
		return Document{}, nil, err
	}
	response.Document.WorkflowRunID = run.ID
	if !published {
		if err = publish(response.Document); err != nil {
			return Document{}, nil, err
		}
	}
	pictures := make([][]byte, len(response.Document.Chapters))
	sem := make(chan struct{}, 3)
	var wg sync.WaitGroup
	var mu sync.Mutex
	var imageErr error
	for i := range pictures {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			select {
			case sem <- struct{}{}:
			case <-ctx.Done():
				return
			}
			defer func() { <-sem }()
			b, err := g.Bridge.Image(ctx, ticket, i)
			mu.Lock()
			defer mu.Unlock()
			if err != nil {
				imageErr = err
			} else {
				pictures[i] = b
			}
		}(i)
	}
	wg.Wait()
	return response.Document, pictures, imageErr
}