package statemachine

import (
	"context"
	"time"

	"go.temporal.io/sdk/workflow"
	"github.com/rockliang/poimen/workflows/action"
	"github.com/rockliang/poimen/workflows/pkg/db"
)

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         []action.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:        []action.EdgeWithWording{},
		Paths:          []QueryPath{},
	}

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

	// Fetch canvas + relations
	var canvasData action.CanvasWithRelationsData
	err := workflow.ExecuteActivity(ctx, action.FetchCanvasRelationsActivity,
		action.FetchCanvasRelationsInput{
			WorkflowID: input.WorkflowID,
			Version:    input.Version,
		},
	).Get(ctx, &canvasData)
	if err != nil {
		return output, err
	}

	// Query Memory System via unified endpoint
	var graphResults action.GraphRAGQueryOutput
	err = workflow.ExecuteActivity(ctx, action.QueryGraphRAGActivity,
		action.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
}