package connector

import (
	"context"
	"fmt"
	"time"

	"github.com/edgeai/gateway/internal/config"
)

// TimeoutManager manages layered timeouts for different phases of request processing.
type TimeoutManager struct {
	cfg *config.TimeoutConfig
}

// NewTimeoutManager creates a new TimeoutManager.
func NewTimeoutManager(cfg *config.TimeoutConfig) *TimeoutManager {
	return &TimeoutManager{cfg: cfg}
}

// TimeoutPhase represents a phase of request processing.
type TimeoutPhase string

const (
	PhaseQueue      TimeoutPhase = "queue"
	PhaseFirstToken TimeoutPhase = "first_token"
	PhaseInference  TimeoutPhase = "inference"
	PhaseTotal      TimeoutPhase = "total"
)

// TimeoutConfig holds resolved timeout values for a specific request.
type TimeoutConfig struct {
	QueueMs      int
	FirstTokenMs int
	InferenceMs  int
	TotalMs      int
	ConnectMs    int
	IdleMs       int
}

// ResolveTimeouts merges request-level timeout overrides with global defaults.
func (tm *TimeoutManager) ResolveTimeouts(reqTimeouts *config.TimeoutConfig, overrides map[string]int) *TimeoutConfig {
	tc := &TimeoutConfig{
		QueueMs:      tm.cfg.DefaultQueueMs,
		FirstTokenMs: tm.cfg.DefaultFirstTokenMs,
		InferenceMs:  tm.cfg.DefaultInferenceMs,
		TotalMs:      tm.cfg.DefaultTotalMs,
		ConnectMs:    tm.cfg.DefaultConnectMs,
		IdleMs:       tm.cfg.DefaultIdleMs,
	}

	if overrides != nil {
		if v, ok := overrides["queue_ms"]; ok && v > 0 {
			tc.QueueMs = v
		}
		if v, ok := overrides["first_token_ms"]; ok && v > 0 {
			tc.FirstTokenMs = v
		}
		if v, ok := overrides["inference_ms"]; ok && v > 0 {
			tc.InferenceMs = v
		}
		if v, ok := overrides["total_ms"]; ok && v > 0 {
			tc.TotalMs = v
		}
	}

	return tc
}

// QueueContext returns a context with the queue timeout.
func (tm *TimeoutManager) QueueContext(parent context.Context, tc *TimeoutConfig) (context.Context, context.CancelFunc) {
	return context.WithTimeout(parent, time.Duration(tc.QueueMs)*time.Millisecond)
}

// InferenceContext returns a context with the inference timeout.
func (tm *TimeoutManager) InferenceContext(parent context.Context, tc *TimeoutConfig) (context.Context, context.CancelFunc) {
	return context.WithTimeout(parent, time.Duration(tc.InferenceMs)*time.Millisecond)
}

// TotalContext returns a context with the total request timeout.
func (tm *TimeoutManager) TotalContext(parent context.Context, tc *TimeoutConfig) (context.Context, context.CancelFunc) {
	return context.WithTimeout(parent, time.Duration(tc.TotalMs)*time.Millisecond)
}

// CheckTimeout returns an error if the given phase has timed out.
func (tm *TimeoutManager) CheckTimeout(phase TimeoutPhase, elapsed time.Duration, tc *TimeoutConfig) error {
	var limit time.Duration
	switch phase {
	case PhaseQueue:
		limit = time.Duration(tc.QueueMs) * time.Millisecond
	case PhaseFirstToken:
		limit = time.Duration(tc.FirstTokenMs) * time.Millisecond
	case PhaseInference:
		limit = time.Duration(tc.InferenceMs) * time.Millisecond
	case PhaseTotal:
		limit = time.Duration(tc.TotalMs) * time.Millisecond
	default:
		return nil
	}

	if elapsed > limit {
		return fmt.Errorf("%s timeout: elapsed %v exceeds limit %v", phase, elapsed, limit)
	}
	return nil
}

// CancelManager manages cancellation propagation from client to inference engine.
type CancelManager struct{}

// NewCancelManager creates a new CancelManager.
func NewCancelManager() *CancelManager {
	return &CancelManager{}
}

// WatchClientDisconnect watches for client connection close and signals cancellation.
// Returns a context that is cancelled when the client disconnects.
func (cm *CancelManager) WatchClientDisconnect(r interface{ Done() <-chan struct{} }, cancel context.CancelFunc) {
	go func() {
		select {
		case <-r.Done():
			cancel()
		}
	}()
}

// PropagateCancel creates a derived context that is cancelled when either the parent
// context is cancelled or the cancel channel is closed.
func (cm *CancelManager) PropagateCancel(parent context.Context, cancelCh <-chan struct{}) (context.Context, context.CancelFunc) {
	ctx, cancel := context.WithCancel(parent)
	go func() {
		select {
		case <-cancelCh:
			cancel()
		case <-ctx.Done():
		}
	}()
	return ctx, cancel
}