package workflow

import (
	"time"

	"go.temporal.io/sdk/temporal"
	"go.temporal.io/sdk/workflow"

	"github.com/rockliang/poimen/workflows/pkg/types"
)

type WorkflowGraphQueryInput struct {
	WorkflowID       string  `json:"workflow_id"`
	Query            string  `json:"query"`
	SearchType       string  `json:"search_type"`
	RelationType     string  `json:"relation_type"`
	Version          int     `json:"version"`
	ConfidenceFloor  float64 `json:"confidence_floor"`
	TopK             int     `json:"top_k"`
	FindPaths        bool    `json:"find_paths"`
	TargetNodeID     string  `json:"target_node_id"`
	MaxPathDepth     int     `json:"max_path_depth"`
	RankingProfile   string  `json:"ranking_profile"`
	IncludeReasoning bool    `json:"include_reasoning"`
}

type WorkflowGraphQueryOutput struct {
	WorkflowID      string                  `json:"workflow_id"`
	Query           string                  `json:"query"`
	Version         int                     `json:"version"`
	ExecutionTimeMs int64                   `json:"execution_time_ms"`
	Results         []types.EdgeWithWording `json:"results"`
	Paths           []QueryPath             `json:"paths"`
	TotalCount      int                     `json:"total_count"`
	HasMore         bool                    `json:"has_more"`
	RankingProfile  string                  `json:"ranking_profile"`
}

type QueryPath struct {
	SourceID   string   `json:"source_id"`
	TargetID   string   `json:"target_id"`
	Distance   int      `json:"distance"`
	PathCount  int      `json:"path_count"`
	NodeIDs    []string `json:"node_ids"`
	Confidence float64  `json:"total_confidence"`
}

func WorkflowGraphQuery(ctx workflow.Context, input WorkflowGraphQueryInput) (WorkflowGraphQueryOutput, error) {
	startTime := time.Now()
	output := WorkflowGraphQueryOutput{
		WorkflowID:     input.WorkflowID,
		Query:          input.Query,
		Version:        input.Version,
		RankingProfile: input.RankingProfile,
		Results:        []types.EdgeWithWording{},
		Paths:          []QueryPath{},
	}

	opts := workflow.ActivityOptions{
		StartToCloseTimeout: 120 * time.Second,
		RetryPolicy: &temporal.RetryPolicy{
			InitialInterval:    2 * time.Second,
			BackoffCoefficient: 2.0,
			MaximumInterval:    10 * time.Second,
			MaximumAttempts:    3,
		},
	}
	ctx = workflow.WithActivityOptions(ctx, opts)

	var canvasData types.CanvasWithRelationsData
	err := workflow.ExecuteActivity(ctx, "FetchCanvasRelationsActivity",
		types.FetchCanvasRelationsInput{
			WorkflowID: input.WorkflowID,
			Version:    input.Version,
		},
	).Get(ctx, &canvasData)
	if err != nil {
		return output, err
	}

	var graphResults types.GraphRAGQueryOutput
	err = workflow.ExecuteActivity(ctx, "QueryGraphRAGActivity",
		types.GraphRAGQueryInput{
			WorkflowID:      input.WorkflowID,
			Query:           input.Query,
			SearchType:      input.SearchType,
			RelationType:    input.RelationType,
			ConfidenceFloor: input.ConfidenceFloor,
			TopK:            input.TopK,
			RankingProfile:  input.RankingProfile,
			Canvas:          canvasData,
		},
	).Get(ctx, &graphResults)
	if err != nil {
		return output, err
	}

	output.Results = graphResults.Edges
	output.TotalCount = graphResults.TotalCount
	output.HasMore = graphResults.HasMore
	output.ExecutionTimeMs = time.Since(startTime).Milliseconds()
	return output, nil
}