package live

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/coder/websocket"
	"github.com/itworx/pulse/internal/auth"
	"github.com/itworx/pulse/internal/correlation"
	"github.com/itworx/pulse/internal/queryplan"
)

const (
	schemaVersion           = 1
	defaultMaxMessage       = 64 << 10
	defaultMaxSubs          = 50
	defaultRateWindow       = time.Minute
	defaultRateMessages     = 120
	defaultHeartbeat        = 15 * time.Second
	defaultIdleTimeout      = 45 * time.Second
	defaultWriteTimeout     = 5 * time.Second
	maxSubscriptionID       = 128
	maxDetailLength         = 500
	defaultOutboundMessages = 32
)

var (
	ErrMessageTooLarge = errors.New("live message exceeds size limit")
	ErrRateLimited     = errors.New("live message rate limit exceeded")
)

type Incoming struct {
	SchemaVersion   int             `json:"schemaVersion"`
	Type            string          `json:"type"`
	SubscriptionID  string          `json:"subscriptionId,omitempty"`
	Query           json.RawMessage `json:"query,omitempty"`
	IntervalSeconds int             `json:"intervalSeconds,omitempty"`
	Nonce           string          `json:"nonce,omitempty"`
}

type Sample struct {
	Series    string            `json:"series"`
	Timestamp time.Time         `json:"timestamp"`
	Value     *float64          `json:"value"`
	Freshness string            `json:"freshness"`
	Labels    map[string]string `json:"labels,omitempty"`
}

type SamplesMessage struct {
	SchemaVersion  int       `json:"schemaVersion"`
	Type           string    `json:"type"`
	SubscriptionID string    `json:"subscriptionId"`
	Sequence       uint64    `json:"sequence"`
	ServerTime     time.Time `json:"serverTime"`
	Samples        []Sample  `json:"samples"`
	Coalesced      bool      `json:"coalesced,omitempty"`
}

type StatusMessage struct {
	SchemaVersion  int    `json:"schemaVersion"`
	Type           string `json:"type"`
	SubscriptionID string `json:"subscriptionId"`
	State          string `json:"state"`
	Detail         string `json:"detail,omitempty"`
}

type ErrorMessage struct {
	SchemaVersion  int    `json:"schemaVersion"`
	Type           string `json:"type"`
	SubscriptionID string `json:"subscriptionId,omitempty"`
	Code           string `json:"code"`
	Message        string `json:"message"`
	CorrelationID  string `json:"correlationId"`
}

type PongMessage struct {
	SchemaVersion int    `json:"schemaVersion"`
	Type          string `json:"type"`
	Nonce         string `json:"nonce"`
}

type Sampler interface {
	Sample(context.Context, queryplan.Request) ([]Sample, error)
}

type Handler struct {
	Planner           *queryplan.Planner
	Sampler           Sampler
	AllowedOrigins    []string
	MaxSubscriptions  int
	MaxMessages       int
	RateWindow        time.Duration
	HeartbeatInterval time.Duration
	IdleTimeout       time.Duration
	WriteTimeout      time.Duration
	Now               func() time.Time
	Registry          *Registry
}

type subscription struct {
	id       string
	request  queryplan.Request
	interval time.Duration
	nextAt   time.Time
	sequence uint64
	lease    *Lease
}

type session struct {
	conn          *websocket.Conn
	principal     auth.Principal
	planner       *queryplan.Planner
	sampler       Sampler
	maxSubs       int
	maxMessages   int
	rateWindow    time.Duration
	heartbeat     time.Duration
	idleTimeout   time.Duration
	writeTimeout  time.Duration
	registry      *Registry
	now           func() time.Time
	correlationID string
	mu            sync.Mutex
	subscriptions map[string]*subscription
	rateStarted   time.Time
	rateCount     int
	outbound      chan []byte
}

func (h Handler) ServeHTTP(response http.ResponseWriter, request *http.Request) {
	principal, ok := auth.PrincipalFromContext(request.Context())
	if !ok {
		http.Error(response, "unauthorized", http.StatusUnauthorized)
		return
	}
	if !auth.Allows(principal.Role, auth.PermissionView) {
		http.Error(response, "forbidden", http.StatusForbidden)
		return
	}
	conn, err := websocket.Accept(response, request, &websocket.AcceptOptions{OriginPatterns: h.AllowedOrigins})
	if err != nil {
		return
	}
	conn.SetReadLimit(int64(defaultMaxMessage))
	registry := h.Registry
	if registry == nil {
		registry = NewRegistry(h.Sampler, RegistryOptions{})
	}
	// Keep the upgraded transport attached to the authenticated request. The
	// session middleware cancels this context on logout and at the finite
	// absolute session deadline; server/request cancellation must also release
	// subscriptions and the socket.
	ctx, cancel := context.WithCancel(request.Context())
	correlationID := correlation.FromContext(request.Context())
	if correlationID == "" {
		correlationID = correlation.New()
	}
	s := &session{
		conn: conn, principal: principal, planner: h.Planner, sampler: h.Sampler,
		maxSubs:      positiveOr(h.MaxSubscriptions, defaultMaxSubs),
		maxMessages:  positiveOr(h.MaxMessages, defaultRateMessages),
		rateWindow:   durationOr(h.RateWindow, defaultRateWindow),
		heartbeat:    durationOr(h.HeartbeatInterval, defaultHeartbeat),
		idleTimeout:  durationOr(h.IdleTimeout, defaultIdleTimeout),
		writeTimeout: durationOr(h.WriteTimeout, defaultWriteTimeout),
		registry:     registry,
		now:          h.Now, correlationID: correlationID, subscriptions: make(map[string]*subscription), outbound: make(chan []byte, defaultOutboundMessages),
	}
	if s.now == nil {
		s.now = func() time.Time { return time.Now().UTC() }
	}
	defer func() {
		cancel()
		s.releaseAll()
		conn.CloseNow()
	}()
	go s.writeLoop(ctx, cancel)
	go s.heartbeatLoop(ctx, cancel)
	go s.sampleLoop(ctx, cancel)
	s.readLoop(ctx)
}

func positiveOr(value, fallback int) int {
	if value > 0 {
		return value
	}
	return fallback
}

func durationOr(value, fallback time.Duration) time.Duration {
	if value > 0 {
		return value
	}
	return fallback
}

func (s *session) readLoop(ctx context.Context) {
	for {
		readCtx, cancel := context.WithTimeout(ctx, s.idleTimeout)
		_, payload, err := s.conn.Read(readCtx)
		cancel()
		if err != nil {
			return
		}
		if err := s.acceptMessage(ctx, payload); err != nil {
			code, detail := "LIVE_MESSAGE_INVALID", "Live bericht is ongeldig."
			if errors.Is(err, ErrMessageTooLarge) {
				code, detail = "LIVE_MESSAGE_LIMIT", "Live bericht overschrijdt de maximale berichtgrootte."
			} else if errors.Is(err, ErrRateLimited) {
				code, detail = "LIVE_RATE_LIMIT", "Te veel live berichten; verbind opnieuw na een korte pauze."
			}
			_ = s.send(ErrorMessage{SchemaVersion: schemaVersion, Type: "error", Code: code, Message: detail, CorrelationID: s.correlationID})
			return
		}
	}
}

func (s *session) acceptMessage(ctx context.Context, payload []byte) error {
	if len(payload) > defaultMaxMessage {
		return ErrMessageTooLarge
	}
	s.mu.Lock()
	now := s.now()
	if s.rateStarted.IsZero() || now.Sub(s.rateStarted) >= s.rateWindow {
		s.rateStarted = now
		s.rateCount = 0
	}
	s.rateCount++
	limited := s.rateCount > s.maxMessages
	s.mu.Unlock()
	if limited {
		return ErrRateLimited
	}
	var message Incoming
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&message); err != nil {
		return err
	}
	if err := decoder.Decode(&struct{}{}); err != io.EOF {
		return errors.New("live message has trailing data")
	}
	if message.SchemaVersion != schemaVersion || !validTypeFields(payload, message.Type) {
		return errors.New("live message schema is invalid")
	}
	switch message.Type {
	case "subscribe":
		return s.subscribe(ctx, message)
	case "unsubscribe":
		return s.unsubscribe(message.SubscriptionID)
	case "ping":
		if !validNonce(message.Nonce) {
			return errors.New("ping nonce is invalid")
		}
		return s.send(PongMessage{SchemaVersion: schemaVersion, Type: "pong", Nonce: message.Nonce})
	case "pong":
		if !validNonce(message.Nonce) {
			return errors.New("pong nonce is invalid")
		}
		return nil
	default:
		return errors.New("live message type is unsupported")
	}
}

func validTypeFields(payload []byte, typ string) bool {
	var fields map[string]json.RawMessage
	if json.Unmarshal(payload, &fields) != nil {
		return false
	}
	allowed := map[string]map[string]struct{}{
		"subscribe":   {"schemaVersion": {}, "type": {}, "subscriptionId": {}, "query": {}, "intervalSeconds": {}},
		"unsubscribe": {"schemaVersion": {}, "type": {}, "subscriptionId": {}},
		"ping":        {"schemaVersion": {}, "type": {}, "nonce": {}},
		"pong":        {"schemaVersion": {}, "type": {}, "nonce": {}},
	}
	known, ok := allowed[typ]
	if !ok {
		return false
	}
	for key := range fields {
		if _, ok := known[key]; !ok {
			return false
		}
	}
	return true
}

func validSubscriptionID(id string) bool {
	return id != "" && utf8.RuneCountInString(id) <= maxSubscriptionID
}

func validNonce(nonce string) bool {
	return utf8.RuneCountInString(nonce) <= maxSubscriptionID
}

func (s *session) subscribe(ctx context.Context, message Incoming) error {
	if !validSubscriptionID(message.SubscriptionID) || len(message.Query) == 0 || message.IntervalSeconds < 1 || message.IntervalSeconds > 300 {
		return s.sendError(message.SubscriptionID, "LIVE_SUBSCRIPTION_INVALID", "Abonnement, query of interval is ongeldig.")
	}
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(message.Query, &fields); err != nil || len(fields) > 20 {
		return s.sendError(message.SubscriptionID, "LIVE_QUERY_INVALID", "Live query moet een begrensd JSON-object zijn.")
	}
	var request queryplan.Request
	decoder := json.NewDecoder(bytes.NewReader(message.Query))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&request); err != nil {
		return s.sendError(message.SubscriptionID, "LIVE_QUERY_INVALID", "Live query bevat onbekende of ongeldige velden.")
	}
	if s.planner == nil {
		return s.sendError(message.SubscriptionID, "LIVE_QUERY_UNAVAILABLE", "Live queryvalidatie is tijdelijk niet beschikbaar.")
	}
	plan, err := s.planner.Plan(auth.WithPrincipal(ctx, s.principal), request)
	if err != nil {
		return s.sendError(message.SubscriptionID, "LIVE_SUBSCRIPTION_DENIED", queryErrorDetail(err))
	}
	now := s.now()
	interval := time.Duration(message.IntervalSeconds) * time.Second
	s.mu.Lock()
	duplicate := false
	sessionLimit := len(s.subscriptions) >= s.maxSubs
	if _, exists := s.subscriptions[message.SubscriptionID]; exists {
		duplicate = true
	}
	s.mu.Unlock()
	if duplicate {
		return s.sendError(message.SubscriptionID, "LIVE_SUBSCRIPTION_CONFLICT", "Subscription-ID bestaat al.")
	}
	if sessionLimit {
		return s.sendError(message.SubscriptionID, "LIVE_SUBSCRIPTION_LIMIT", "Maximum aantal live abonnementen bereikt.")
	}
	lease, err := s.registry.Acquire(plan.Request, interval)
	if err != nil {
		if errors.Is(err, ErrRegistryLimit) {
			return s.sendError(message.SubscriptionID, "LIVE_REGISTRY_LIMIT", "Maximum aantal gedeelde live queries bereikt.")
		}
		return s.sendError(message.SubscriptionID, "LIVE_QUERY_INVALID", "Live query kon niet worden genormaliseerd.")
	}
	s.mu.Lock()
	if _, exists := s.subscriptions[message.SubscriptionID]; exists || len(s.subscriptions) >= s.maxSubs {
		s.mu.Unlock()
		lease.Release()
		return s.sendError(message.SubscriptionID, "LIVE_SUBSCRIPTION_CONFLICT", "Subscription-ID bestaat al.")
	}
	s.subscriptions[message.SubscriptionID] = &subscription{id: message.SubscriptionID, request: plan.Request, interval: interval, nextAt: now.Add(interval), lease: lease}
	s.mu.Unlock()
	return s.send(StatusMessage{SchemaVersion: schemaVersion, Type: "status", SubscriptionID: message.SubscriptionID, State: "subscribed"})
}

func queryErrorDetail(err error) string {
	var plannerError queryplan.Error
	if errors.As(err, &plannerError) {
		return plannerError.Detail
	}
	return "Live query kon niet worden geautoriseerd."
}

func (s *session) unsubscribe(id string) error {
	if !validSubscriptionID(id) {
		return s.sendError(id, "LIVE_SUBSCRIPTION_INVALID", "Subscription-ID is ongeldig.")
	}
	s.mu.Lock()
	sub, exists := s.subscriptions[id]
	delete(s.subscriptions, id)
	s.mu.Unlock()
	if exists && sub.lease != nil {
		sub.lease.Release()
	}
	return s.send(StatusMessage{SchemaVersion: schemaVersion, Type: "status", SubscriptionID: id, State: "unsubscribed"})
}

func (s *session) releaseAll() {
	s.mu.Lock()
	subscriptions := make([]*subscription, 0, len(s.subscriptions))
	for _, sub := range s.subscriptions {
		subscriptions = append(subscriptions, sub)
	}
	s.subscriptions = make(map[string]*subscription)
	s.mu.Unlock()
	for _, sub := range subscriptions {
		if sub.lease != nil {
			sub.lease.Release()
		}
	}
}

func (s *session) sendError(subscriptionID, code, message string) error {
	if len(message) > maxDetailLength {
		message = message[:maxDetailLength]
	}
	return s.send(ErrorMessage{SchemaVersion: schemaVersion, Type: "error", SubscriptionID: subscriptionID, Code: code, Message: message, CorrelationID: s.correlationID})
}

func (s *session) send(message any) error {
	payload, err := json.Marshal(message)
	if err != nil {
		return err
	}
	if len(payload) > defaultMaxMessage {
		return ErrMessageTooLarge
	}
	sample := false
	switch message.(type) {
	case SamplesMessage, *SamplesMessage:
		sample = true
	}
	if !sample {
		ctx, cancel := context.WithTimeout(context.Background(), s.writeTimeout)
		defer cancel()
		return s.conn.Write(ctx, websocket.MessageText, payload)
	}
	if s.outbound == nil {
		return ErrOutboundBackpressure
	}
	select {
	case s.outbound <- payload:
		return nil
	default:
		return nil
	}
}

func (s *session) writeLoop(ctx context.Context, cancel context.CancelFunc) {
	for {
		select {
		case payload := <-s.outbound:
			writeCtx, writeCancel := context.WithTimeout(ctx, s.writeTimeout)
			err := s.conn.Write(writeCtx, websocket.MessageText, payload)
			writeCancel()
			if err != nil {
				cancel()
				return
			}
		case <-ctx.Done():
			return
		}
	}
}

func (s *session) heartbeatLoop(ctx context.Context, cancel context.CancelFunc) {
	ticker := time.NewTicker(s.heartbeat)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			pingCtx, pingCancel := context.WithTimeout(ctx, s.writeTimeout)
			err := s.conn.Ping(pingCtx)
			pingCancel()
			if err != nil {
				cancel()
				return
			}
		case <-ctx.Done():
			return
		}
	}
}

func (s *session) sampleLoop(ctx context.Context, cancel context.CancelFunc) {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			now := s.now()
			s.mu.Lock()
			due := make([]*subscription, 0, len(s.subscriptions))
			for _, sub := range s.subscriptions {
				if !now.Before(sub.nextAt) {
					sub.nextAt = now.Add(sub.interval)
					due = append(due, sub)
				}
			}
			s.mu.Unlock()
			for _, sub := range due {
				samples, err := s.registry.Sample(ctx, sub.request)
				if err != nil {
					_ = s.sendError(sub.id, "LIVE_SAMPLE_UNAVAILABLE", "Live sample is tijdelijk niet beschikbaar.")
					continue
				}
				s.mu.Lock()
				current, exists := s.subscriptions[sub.id]
				if !exists {
					s.mu.Unlock()
					continue
				}
				current.sequence++
				sequence := current.sequence
				s.mu.Unlock()
				if len(samples) > 10000 {
					_ = s.sendError(sub.id, "LIVE_SAMPLE_LIMIT", "Live sample bevat te veel punten.")
					continue
				}
				if err := s.send(SamplesMessage{SchemaVersion: schemaVersion, Type: "samples", SubscriptionID: sub.id, Sequence: sequence, ServerTime: now.UTC(), Samples: samples}); err != nil {
					cancel()
					return
				}
			}
		case <-ctx.Done():
			return
		}
	}
}