package agent

import (
	"context"
	"math"

	"github.com/tidbcloud/AgentTeam/kits/agentteam-insight/internal/graphbuilder/batching"
	"github.com/tidbcloud/AgentTeam/kits/agentteam-insight/internal/graphbuilder/extractor"
	"github.com/tidbcloud/AgentTeam/kits/agentteam-insight/internal/graphbuilder/scanner"
	"github.com/tidbcloud/AgentTeam/kits/agentteam-insight/internal/graphbuilder/schema"
)

// Request describes one optional graph enhancement pass. The orchestrator
// dispatches one request per batch so a file-analyzer runner receives
// the same batchIndex, pre-resolved imports, and cross-batch neighbor context
// as the graph builder batch plan.
type Request struct {
	OutputLanguage  string                         `json:"outputLanguage"`
	LocaleAddendum  string                         `json:"localeAddendum,omitempty"`
	Project         ProjectContext                 `json:"project"`
	BatchIndex      int                            `json:"batchIndex"`
	TotalBatches    int                            `json:"totalBatches"`
	Files           []scanner.FileInfo             `json:"files"`
	BatchImportData map[string][]string            `json:"batchImportData"`
	NeighborMap     map[string][]batching.Neighbor `json:"neighborMap"`
	Analyses        []extractor.StructuralAnalysis `json:"-"`
}

type ProjectContext struct {
	Name                     string   `json:"name"`
	Description              string   `json:"description"`
	ReadmeHead               string   `json:"readmeHead"`
	Languages                []string `json:"languages"`
	Frameworks               []string `json:"frameworks"`
	SupportedImportResolvers []string `json:"supportedImportResolvers"`
}

// Runner enhances a graph without owning persistence or orchestration.
type Runner interface {
	Enhance(context.Context, schema.KnowledgeGraph, Request) (schema.KnowledgeGraph, error)
}

// RunnerFunc adapts a function into a Runner for tests and local seams.
type RunnerFunc func(context.Context, schema.KnowledgeGraph, Request) (schema.KnowledgeGraph, error)

func (fn RunnerFunc) Enhance(ctx context.Context, graph schema.KnowledgeGraph, request Request) (schema.KnowledgeGraph, error) {
	return fn(ctx, graph, request)
}

// Options controls the optional enhancement pass.
type Options struct {
	Enabled        bool
	Runner         Runner
	OutputLanguage string
}

// Enhance runs the optional agent seam and validates the returned graph before
// handing it back to the orchestrator. Disabled or unconfigured seams are a
// deterministic no-op, so default CLI runs never depend on an LLM.
func Enhance(ctx context.Context, graph schema.KnowledgeGraph, options Options) (schema.KnowledgeGraph, bool, error) {
	if !options.Enabled || options.Runner == nil {
		return graph, false, nil
	}
	if err := ctx.Err(); err != nil {
		return schema.KnowledgeGraph{}, false, err
	}
	enhanced, err := options.Runner.Enhance(ctx, graph, Request{OutputLanguage: options.OutputLanguage})
	if err != nil {
		return schema.KnowledgeGraph{}, false, err
	}
	result := schema.ValidateGraph(enhanced)
	if !result.Success || result.Data == nil {
		return schema.KnowledgeGraph{}, false, &ValidationError{Message: firstNonEmpty(result.Fatal, "invalid enhanced graph")}
	}
	return *result.Data, true, nil
}

// EnhanceBatch runs the optional agent seam for a single file-analysis batch.
// Disabled or nil-runner paths are deterministic no-ops but still return the
// same contract shape the orchestrator would pass to a real runner.
func EnhanceBatch(ctx context.Context, graph schema.KnowledgeGraph, options Options, request Request) (schema.KnowledgeGraph, bool, error) {
	if !options.Enabled || options.Runner == nil {
		return graph, false, nil
	}
	if err := ctx.Err(); err != nil {
		return schema.KnowledgeGraph{}, false, err
	}
	request.OutputLanguage = options.OutputLanguage
	enhanced, err := options.Runner.Enhance(ctx, graph, request)
	if err != nil {
		return schema.KnowledgeGraph{}, false, err
	}
	// Do not run full graph validation on a single batch fragment here:
	// legitimate file-analyzer output may include cross-batch edges whose
	// target nodes are supplied by another batch. The orchestrator validates
	// the assembled graph after all fragments have been merged.
	return defaultBatchEdgeFields(enhanced), true, nil
}

func defaultBatchEdgeFields(graph schema.KnowledgeGraph) schema.KnowledgeGraph {
	for i := range graph.Edges {
		if graph.Edges[i].Direction == "" {
			graph.Edges[i].Direction = "forward"
		}
		if math.IsNaN(graph.Edges[i].Weight) {
			graph.Edges[i].Weight = 0.5
		}
	}
	return graph
}

// ValidationError reports an invalid graph returned by the optional runner.
type ValidationError struct {
	Message string
}

func (err *ValidationError) Error() string {
	return "agent enhancement returned invalid graph: " + err.Message
}

func firstNonEmpty(values ...string) string {
	for _, value := range values {
		if value != "" {
			return value
		}
	}
	return ""
}