package graphiti

import (
	"context"
	"fmt"
	"time"

	graphiti "github.com/vxcontrol/graphiti-go-client"
)

// Re-export types from the graphiti-go-client package for convenience
type (
	Observation        = graphiti.Observation
	Message            = graphiti.Message
	AddMessagesRequest = graphiti.AddMessagesRequest

	// Search request/response types
	TemporalSearchRequest            = graphiti.TemporalSearchRequest
	TemporalSearchResponse           = graphiti.TemporalSearchResponse
	EntityRelationshipSearchRequest  = graphiti.EntityRelationshipSearchRequest
	EntityRelationshipSearchResponse = graphiti.EntityRelationshipSearchResponse
	DiverseSearchRequest             = graphiti.DiverseSearchRequest
	DiverseSearchResponse            = graphiti.DiverseSearchResponse
	EpisodeContextSearchRequest      = graphiti.EpisodeContextSearchRequest
	EpisodeContextSearchResponse     = graphiti.EpisodeContextSearchResponse
	SuccessfulToolsSearchRequest     = graphiti.SuccessfulToolsSearchRequest
	SuccessfulToolsSearchResponse    = graphiti.SuccessfulToolsSearchResponse
	RecentContextSearchRequest       = graphiti.RecentContextSearchRequest
	RecentContextSearchResponse      = graphiti.RecentContextSearchResponse
	EntityByLabelSearchRequest       = graphiti.EntityByLabelSearchRequest
	EntityByLabelSearchResponse      = graphiti.EntityByLabelSearchResponse

	// Common types used in search responses
	NodeResult      = graphiti.NodeResult
	EdgeResult      = graphiti.EdgeResult
	EpisodeResult   = graphiti.EpisodeResult
	CommunityResult = graphiti.CommunityResult
	TimeWindow      = graphiti.TimeWindow
)

// Client wraps the Graphiti client with Pentagi-specific functionality
type Client struct {
	client  *graphiti.Client
	enabled bool
	timeout time.Duration
}

// The verdict is permanent for the process: on failure the caller keeps a disabled client for its whole
// lifetime, so one timeout against a slow endpoint must not decide it.
const (
	healthCheckAttempts = 3
	healthCheckBackoff  = 2 * time.Second
)

// NewClient creates a new Graphiti client wrapper
func NewClient(url string, timeout time.Duration, enabled bool) (*Client, error) {
	if !enabled {
		return &Client{enabled: false}, nil
	}

	client := graphiti.NewClient(url, graphiti.WithTimeout(timeout))

	var err error
	for attempt := 1; ; attempt++ {
		if _, err = client.HealthCheck(); err == nil {
			break
		}

		if attempt == healthCheckAttempts {
			return nil, fmt.Errorf("graphiti health check failed after %d attempts: %w", attempt, err)
		}

		time.Sleep(healthCheckBackoff)
	}

	return &Client{
		client:  client,
		enabled: true,
		timeout: timeout,
	}, nil
}

// IsEnabled returns whether Graphiti integration is active
func (c *Client) IsEnabled() bool {
	return c != nil && c.enabled
}

// GetTimeout returns the configured timeout duration
func (c *Client) GetTimeout() time.Duration {
	if c == nil {
		return 0
	}
	return c.timeout
}

// awaitCall hands the caller back the moment ctx ends. The call it wraps keeps
// running to its own timeout, so this bounds the wait, not the work.
func awaitCall[T any](ctx context.Context, call func() (T, error)) (T, error) {
	type outcome struct {
		value T
		err   error
	}

	done := make(chan outcome, 1)
	go func() {
		value, err := call()
		done <- outcome{value: value, err: err}
	}()

	select {
	case result := <-done:
		return result.value, result.err
	case <-ctx.Done():
		var zero T

		return zero, ctx.Err()
	}
}

// AddMessages adds messages to Graphiti (no-op if disabled)
func (c *Client) AddMessages(ctx context.Context, req graphiti.AddMessagesRequest) error {
	if !c.IsEnabled() {
		return nil
	}

	_, err := awaitCall(ctx, func() (*graphiti.Result, error) { return c.client.AddMessages(req) })

	return err
}

// TemporalWindowSearch searches within a time window
func (c *Client) TemporalWindowSearch(ctx context.Context, req TemporalSearchRequest) (*TemporalSearchResponse, error) {
	if !c.IsEnabled() {
		return nil, fmt.Errorf("graphiti is not enabled")
	}
	return awaitCall(ctx, func() (*TemporalSearchResponse, error) { return c.client.TemporalWindowSearch(req) })
}

// EntityRelationshipsSearch finds relationships from a center node
func (c *Client) EntityRelationshipsSearch(ctx context.Context, req EntityRelationshipSearchRequest) (*EntityRelationshipSearchResponse, error) {
	if !c.IsEnabled() {
		return nil, fmt.Errorf("graphiti is not enabled")
	}
	return awaitCall(ctx, func() (*EntityRelationshipSearchResponse, error) { return c.client.EntityRelationshipsSearch(req) })
}

// DiverseResultsSearch gets diverse, non-redundant results
func (c *Client) DiverseResultsSearch(ctx context.Context, req DiverseSearchRequest) (*DiverseSearchResponse, error) {
	if !c.IsEnabled() {
		return nil, fmt.Errorf("graphiti is not enabled")
	}
	return awaitCall(ctx, func() (*DiverseSearchResponse, error) { return c.client.DiverseResultsSearch(req) })
}

// EpisodeContextSearch searches through agent responses and tool execution records
func (c *Client) EpisodeContextSearch(ctx context.Context, req EpisodeContextSearchRequest) (*EpisodeContextSearchResponse, error) {
	if !c.IsEnabled() {
		return nil, fmt.Errorf("graphiti is not enabled")
	}
	return awaitCall(ctx, func() (*EpisodeContextSearchResponse, error) { return c.client.EpisodeContextSearch(req) })
}

// SuccessfulToolsSearch finds successful tool executions and attack patterns
func (c *Client) SuccessfulToolsSearch(ctx context.Context, req SuccessfulToolsSearchRequest) (*SuccessfulToolsSearchResponse, error) {
	if !c.IsEnabled() {
		return nil, fmt.Errorf("graphiti is not enabled")
	}
	return awaitCall(ctx, func() (*SuccessfulToolsSearchResponse, error) { return c.client.SuccessfulToolsSearch(req) })
}

// RecentContextSearch retrieves recent relevant context
func (c *Client) RecentContextSearch(ctx context.Context, req RecentContextSearchRequest) (*RecentContextSearchResponse, error) {
	if !c.IsEnabled() {
		return nil, fmt.Errorf("graphiti is not enabled")
	}
	return awaitCall(ctx, func() (*RecentContextSearchResponse, error) { return c.client.RecentContextSearch(req) })
}

// EntityByLabelSearch searches for entities by label/type
func (c *Client) EntityByLabelSearch(ctx context.Context, req EntityByLabelSearchRequest) (*EntityByLabelSearchResponse, error) {
	if !c.IsEnabled() {
		return nil, fmt.Errorf("graphiti is not enabled")
	}
	return awaitCall(ctx, func() (*EntityByLabelSearchResponse, error) { return c.client.EntityByLabelSearch(req) })
}