package wire

import (
	"context"
	"encoding/json"
	"log/slog"
	"strings"
	"time"
)

func (c *Protocol) trackElasticRequestLink(link *elasticRequestLink) {
	if link == nil || strings.TrimSpace(link.linkID) == "" {
		return
	}
	var old *elasticRequestLink
	c.elasticLinksMu.Lock()
	if c.elasticRequestLinks == nil {
		c.elasticRequestLinks = make(map[string]*elasticRequestLink)
	}
	old = c.elasticRequestLinks[link.linkID]
	c.elasticRequestLinks[link.linkID] = link
	c.elasticLinksMu.Unlock()
	if old != nil {
		old.close()
	}
}

func (c *Protocol) untrackElasticRequestLink(link *elasticRequestLink) {
	if link == nil {
		return
	}
	c.elasticLinksMu.Lock()
	if c.elasticRequestLinks[strings.TrimSpace(link.linkID)] == link {
		delete(c.elasticRequestLinks, strings.TrimSpace(link.linkID))
	}
	c.elasticLinksMu.Unlock()
}

func (c *Protocol) closeElasticRequestLinks(reason string) {
	c.elasticLinksMu.Lock()
	links := make([]*elasticRequestLink, 0, len(c.elasticRequestLinks))
	for _, link := range c.elasticRequestLinks {
		links = append(links, link)
	}
	c.elasticRequestLinks = make(map[string]*elasticRequestLink)
	c.elasticLinksMu.Unlock()
	for _, link := range links {
		link.close()
	}
	if strings.TrimSpace(reason) != "" && len(links) > 0 {
		slog.Info("AgentTeam: closed Elastic Request Links", "channel", c.id, "reason", reason, "count", len(links))
	}
}

func (c *Protocol) runElasticRequestLink(ctx context.Context, link *elasticRequestLink, open elasticRequestLinkOpenMsg) {
	idleTimeout := durationFromMillis(open.IdleTimeoutMs, defaultElasticRequestLinkIdleMs)
	maxAge := durationFromMillis(open.MaxAgeMs, defaultElasticRequestLinkMaxAgeMs)
	idleTimer := time.NewTimer(idleTimeout)
	defer idleTimer.Stop()
	maxAgeTimer := time.NewTimer(maxAge)
	defer maxAgeTimer.Stop()

	readCh := make(chan []byte, 1)
	errCh := make(chan error, 1)
	readerDone := make(chan struct{})
	go func() {
		defer close(readerDone)
		for {
			_, data, err := link.conn.ReadMessage()
			if err != nil {
				errCh <- err
				return
			}
			select {
			case readCh <- data:
			case <-ctx.Done():
				return
			}
		}
	}()
	defer func() {
		// Closing the link cancels ctx and unblocks both a carrier ReadMessage and
		// a reader waiting to publish into readCh. Join that exact reader before the
		// link owner retires from the connection-run registry.
		link.close()
		<-readerDone
		c.untrackElasticRequestLink(link)
		close(link.done)
	}()

	for {
		select {
		case <-ctx.Done():
			return
		case <-maxAgeTimer.C:
			if link.isBusy() {
				link.requestDrain()
				resetTimer(maxAgeTimer, idleTimeout)
				continue
			}
			return
		case <-idleTimer.C:
			if !link.isBusy() {
				return
			}
			resetTimer(idleTimer, idleTimeout)
		case err := <-errCh:
			if err != nil && !isNormalClose(err) {
				slog.Debug("AgentTeam: Elastic Request Link closed", "channel", c.id, "linkId", link.linkID, "reason", "read_failed")
			}
			return
		case data := <-readCh:
			resetTimer(idleTimer, idleTimeout)
			c.dispatchElasticRequestLinkMessage(ctx, link, data)
		}
	}
}

func resetTimer(timer *time.Timer, duration time.Duration) {
	if !timer.Stop() {
		select {
		case <-timer.C:
		default:
		}
	}
	timer.Reset(duration)
}

func durationFromMillis(value int64, fallback int64) time.Duration {
	if value <= 0 {
		value = fallback
	}
	return time.Duration(value) * time.Millisecond
}

func (c *Protocol) dispatchElasticRequestLinkMessage(ctx context.Context, link *elasticRequestLink, raw []byte) {
	var env envelope
	if err := json.Unmarshal(raw, &env); err != nil {
		slog.Debug("AgentTeam: Elastic Request Link bad message", "channel", c.id, "linkId", link.linkID, "error", err)
		return
	}
	switch env.Type {
	case remoteRequestAssignType:
		c.handleRemoteRequestAssignment(ctx, link, raw)
	case elasticRequestLinkDrainType:
		if link.requestDrain() {
			link.close()
		}
	case elasticRequestLinkCancelType:
		var msg elasticRequestLinkCancelMsg
		if err := json.Unmarshal(raw, &msg); err == nil {
			link.cancelAssignment(msg.RemoteRequestID, msg.CancelToken)
		}
	default:
		slog.Debug("AgentTeam: Elastic Request Link unhandled message type", "channel", c.id, "linkId", link.linkID, "type", env.Type)
	}
}

func (c *Protocol) handleElasticRequestLinkDrain(raw []byte) {
	var msg elasticRequestLinkDrainMsg
	if err := json.Unmarshal(raw, &msg); err != nil {
		slog.Debug("AgentTeam: Elastic Request Link drain parse error", "channel", c.id, "error", err)
		return
	}
	if link := c.elasticRequestLink(msg.LinkID); link != nil {
		if link.requestDrain() {
			link.close()
		}
	}
}

func (c *Protocol) handleElasticRequestLinkCancel(raw []byte) {
	var msg elasticRequestLinkCancelMsg
	if err := json.Unmarshal(raw, &msg); err != nil {
		slog.Debug("AgentTeam: Elastic Request Link cancel parse error", "channel", c.id, "error", err)
		return
	}
	if link := c.elasticRequestLink(msg.LinkID); link != nil {
		link.cancelAssignment(msg.RemoteRequestID, msg.CancelToken)
	}
}

func (c *Protocol) elasticRequestLink(linkID string) *elasticRequestLink {
	c.elasticLinksMu.Lock()
	defer c.elasticLinksMu.Unlock()
	return c.elasticRequestLinks[strings.TrimSpace(linkID)]
}