package providerhealth

import (
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"time"

	"aigw/internal/catalog"
	"aigw/internal/domain"
)

type ProbeMetrics interface {
	ProviderProbe(success bool)
}

type ProbeOptions struct {
	Enabled  bool
	Interval time.Duration
	Timeout  time.Duration
	Catalog  *catalog.Catalog
	Tracker  *Tracker
	Metrics  ProbeMetrics
	Logger   *slog.Logger
	Client   *http.Client
}

type Prober struct {
	enabled  bool
	interval time.Duration
	timeout  time.Duration
	catalog  *catalog.Catalog
	tracker  *Tracker
	metrics  ProbeMetrics
	logger   *slog.Logger
	client   *http.Client
}

type probeTarget struct {
	provider domain.Provider
	keys     []RouteKey
}

func NewProber(options ProbeOptions) *Prober {
	if options.Interval <= 0 {
		options.Interval = 30 * time.Second
	}
	if options.Timeout <= 0 {
		options.Timeout = 5 * time.Second
	}
	if options.Logger == nil {
		options.Logger = slog.Default()
	}
	if options.Client == nil {
		transport := http.DefaultTransport.(*http.Transport).Clone()
		transport.Proxy = http.ProxyFromEnvironment
		options.Client = &http.Client{Transport: transport, Timeout: options.Timeout}
	}
	return &Prober{enabled: options.Enabled, interval: options.Interval, timeout: options.Timeout,
		catalog: options.Catalog, tracker: options.Tracker, metrics: options.Metrics, logger: options.Logger, client: options.Client}
}

func (p *Prober) Run(ctx context.Context) {
	if p == nil || !p.enabled || p.catalog == nil || p.tracker == nil {
		return
	}
	p.ProbeOnce(ctx)
	ticker := time.NewTicker(p.interval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			p.ProbeOnce(ctx)
		}
	}
}

func (p *Prober) ProbeOnce(ctx context.Context) {
	for _, target := range p.targets() {
		if ctx.Err() != nil {
			return
		}
		p.probe(ctx, target)
	}
}

func (p *Prober) targets() []probeTarget {
	providerIndexes := make(map[string]int)
	seen := make(map[RouteKey]struct{})
	result := make([]probeTarget, 0)
	for _, model := range p.catalog.AllModels() {
		for _, route := range model.Routes {
			key := RouteKey{ModelID: model.ID, ProviderID: route.Provider.ID, WireAPI: route.Provider.EffectiveWireAPI()}
			if key.ProviderID == "" {
				continue
			}
			if _, exists := seen[key]; exists {
				continue
			}
			seen[key] = struct{}{}
			index, exists := providerIndexes[route.Provider.ID]
			if !exists {
				index = len(result)
				providerIndexes[route.Provider.ID] = index
				result = append(result, probeTarget{provider: route.Provider})
			}
			result[index].keys = append(result[index].keys, key)
		}
	}
	return result
}

func (p *Prober) probe(parent context.Context, target probeTarget) {
	ctx, cancel := context.WithTimeout(parent, p.timeout)
	defer cancel()
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, providerModelsURL(target.provider), nil)
	if err != nil {
		p.observe(target, 0, 0, false)
		return
	}
	request.Header.Set("Accept", "application/json")
	request.Header.Set("User-Agent", "aigw-health/0.1")
	if target.provider.Protocol == domain.ProtocolAnthropic {
		request.Header.Set("x-api-key", target.provider.APIKey)
		request.Header.Set("anthropic-version", "2023-06-01")
	} else {
		request.Header.Set("Authorization", "Bearer "+target.provider.APIKey)
	}
	started := time.Now()
	response, err := p.client.Do(request)
	latency := time.Since(started)
	if err != nil {
		p.observe(target, 0, latency, false)
		p.logger.Warn("provider_probe_failed", "provider_id", target.provider.ID, "error", err)
		return
	}
	body, readErr := io.ReadAll(io.LimitReader(response.Body, 1<<20))
	_ = response.Body.Close()
	success := response.StatusCode >= 200 && response.StatusCode < 300 && readErr == nil &&
		strings.HasPrefix(strings.ToLower(response.Header.Get("Content-Type")), "application/json") && validModelsEnvelope(body)
	p.observe(target, response.StatusCode, latency, success)
	if !success {
		p.logger.Warn("provider_probe_failed", "provider_id", target.provider.ID, "status_code", response.StatusCode)
	}
}

func (p *Prober) observe(target probeTarget, statusCode int, latency time.Duration, success bool) {
	for _, key := range target.keys {
		p.tracker.Observe(key, Observation{StatusCode: statusCode, Latency: latency, Failed: !success, Active: true})
	}
	if p.metrics != nil {
		p.metrics.ProviderProbe(success)
	}
}

func validModelsEnvelope(body []byte) bool {
	var envelope struct {
		Data []json.RawMessage `json:"data"`
	}
	return json.Unmarshal(body, &envelope) == nil && envelope.Data != nil
}

func providerModelsURL(provider domain.Provider) string {
	baseURL := strings.TrimRight(provider.BaseURL, "/")
	return baseURL + "/models"
}