package service

import (
	"context"
	"fmt"

	"gen2d/internal/logger"

	"github.com/cloudwego/eino/compose"
)

const (
	nodePromptOptimizer   = "prompt_optimizer"
	nodeAssetGenerator    = "asset_generator"
	nodeQualitySupervisor = "quality_supervisor"
	nodeFormatAdapter     = "format_adapter"
)

// NewGenerateGraph 创建生成管线 Graph（PromptOptimizer → AssetGenerator → QualitySupervisor → FormatAdapter）。
//
//	START → PromptOptimizer → AssetGenerator → QualitySupervisor
//	                                             ├── pass → FormatAdapter → END
//	                                             └── fail, retry<3 → PromptOptimizer
//	                                             └── fail, retry>=3 → FormatAdapter (降级)
func NewGenerateGraph() (*compose.Graph[PipelineInput, PipelineOutput], error) {
	g := compose.NewGraph[PipelineInput, PipelineOutput](
		compose.WithGenLocalState(func(ctx context.Context) *PipelineState {
			return &PipelineState{}
		}),
	)

	if err := g.AddLambdaNode(nodePromptOptimizer, promptOptimizerNode,
		compose.WithStatePreHandler(promptOptimizerPreHandler),
		compose.WithStatePostHandler(promptOptimizerPostHandler),
	); err != nil {
		return nil, fmt.Errorf("add %s node: %w", nodePromptOptimizer, err)
	}

	if err := g.AddLambdaNode(nodeAssetGenerator, assetGeneratorNode,
		compose.WithStatePostHandler(assetGeneratorPostHandler),
	); err != nil {
		return nil, fmt.Errorf("add %s node: %w", nodeAssetGenerator, err)
	}

	if err := g.AddLambdaNode(nodeQualitySupervisor, qualitySupervisorNode); err != nil {
		return nil, fmt.Errorf("add %s node: %w", nodeQualitySupervisor, err)
	}

	if err := g.AddLambdaNode(nodeFormatAdapter, formatAdapterNode); err != nil {
		return nil, fmt.Errorf("add %s node: %w", nodeFormatAdapter, err)
	}

	// 连线：START → PromptOptimizer → AssetGenerator → Supervisor
	if err := g.AddEdge(compose.START, nodePromptOptimizer); err != nil {
		return nil, fmt.Errorf("add edge START->%s: %w", nodePromptOptimizer, err)
	}
	if err := g.AddEdge(nodePromptOptimizer, nodeAssetGenerator); err != nil {
		return nil, fmt.Errorf("add edge %s->%s: %w", nodePromptOptimizer, nodeAssetGenerator, err)
	}
	if err := g.AddEdge(nodeAssetGenerator, nodeQualitySupervisor); err != nil {
		return nil, fmt.Errorf("add edge %s->%s: %w", nodeAssetGenerator, nodeQualitySupervisor, err)
	}
	if err := g.AddEdge(nodeFormatAdapter, compose.END); err != nil {
		return nil, fmt.Errorf("add edge %s->END: %w", nodeFormatAdapter, err)
	}

	// 连线：质检分支（从 state.NextNode 读取路由目标）
	if err := g.AddBranch(nodeQualitySupervisor, compose.NewGraphBranch(
		func(ctx context.Context, _ PipelineInput) (string, error) {
			var next string
			_ = compose.ProcessState[*PipelineState](ctx, func(_ context.Context, state *PipelineState) error {
				next = state.NextNode
				return nil
			})
			return next, nil
		},
		map[string]bool{nodePromptOptimizer: true, nodeFormatAdapter: true},
	)); err != nil {
		return nil, fmt.Errorf("add branch at %s: %w", nodeQualitySupervisor, err)
	}

	return g, nil
}

// RunPipeline 编译并执行生成管线。
func RunPipeline(ctx context.Context, in PipelineInput) (*PipelineOutput, error) {
	l := logger.FromCtx(ctx)

	l.Info("pipeline started",
		"asset_type", in.AssetType,
		"tags", in.Tags,
		"resolution", in.Params.Resolution,
	)

	g, err := NewGenerateGraph()
	if err != nil {
		l.Error("pipeline create graph failed", "error", err)
		return nil, fmt.Errorf("create graph: %w", err)
	}

	r, err := g.Compile(ctx, compose.WithMaxRunSteps(20))
	if err != nil {
		l.Error("pipeline compile failed", "error", err)
		return nil, fmt.Errorf("compile graph: %w", err)
	}

	output, err := r.Invoke(ctx, in)
	if err != nil {
		l.Error("pipeline invoke failed", "error", err)
		return nil, fmt.Errorf("invoke pipeline: %w", err)
	}

	l.Info("pipeline completed", "asset_count", len(output.Assets))
	return &output, nil
}