package wire

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

const elasticRequestLinkDialTimeout = 5 * time.Second

type elasticRequestLinkDialContext struct {
	Dialer  *websocket.Dialer
	Headers http.Header
}

func (c *Protocol) handleElasticRequestLinkOpen(ctx context.Context, raw []byte) {
	var msg elasticRequestLinkOpenMsg
	if err := json.Unmarshal(raw, &msg); err != nil {
		slog.Debug("AgentTeam: Elastic Request Link open parse error", "channel", c.id, "error", err)
		return
	}
	if err := c.validateElasticRequestLinkOpen(msg); err != nil {
		slog.Warn("AgentTeam: rejected Elastic Request Link open", "channel", c.id, "reason", err.Error())
		return
	}
	c.startConnectionRunTask(ctx, func(ctx context.Context) { c.dialElasticRequestLink(ctx, msg) })
}

func (c *Protocol) validateElasticRequestLinkOpen(msg elasticRequestLinkOpenMsg) error {
	if strings.TrimSpace(msg.LinkID) == "" || strings.TrimSpace(msg.ConnectURL) == "" || strings.TrimSpace(msg.Token) == "" {
		return errors.New("missing_link_grant")
	}
	if msg.ProtocolVersion != daemonTransportSchemaVersion {
		return errors.New("protocol_mismatch")
	}
	if strings.TrimSpace(msg.EndpointRegistryVersion) != defaultEndpointRegistryVersion {
		return errors.New("registry_mismatch")
	}
	registry := c.remoteEndpointRegistry
	if registry == nil {
		registry = defaultDesktopRemoteEndpointRegistry()
	}
	// Validate the Server-supplied endpoint allowlist against the local registry
	// before dialing so Desktop never accepts arbitrary endpoint IDs.
	if len(msg.AllowedEndpoints) == 0 {
		return errors.New("endpoint_not_allowed")
	}
	for _, endpoint := range msg.AllowedEndpoints {
		if _, ok := registry.lookup(endpoint.EndpointID, endpoint.Version); !ok {
			return errors.New("endpoint_not_allowed")
		}
		if strings.TrimSpace(msg.Purpose) == "warm" && !desktopRemoteEndpointAllowsWarmCarrier(endpoint.EndpointID, endpoint.Version) {
			return errors.New("warm_endpoint_not_allowed")
		}
	}
	return nil
}

func (c *Protocol) dialElasticRequestLink(ctx context.Context, msg elasticRequestLinkOpenMsg) {
	serverURL, _, workspaceID, _, _, _ := c.identitySnapshot()
	wsURL, err := buildElasticRequestLinkURL(serverURL, msg.ConnectURL)
	if err != nil {
		slog.Warn("AgentTeam: invalid Elastic Request Link connect URL", "channel", c.id, "linkId", msg.LinkID, "reason", "invalid_connect_url")
		return
	}
	dialCtx, cancel := context.WithTimeout(ctx, elasticRequestLinkDialTimeout)
	defer cancel()

	dialer := websocket.Dialer{HandshakeTimeout: elasticRequestLinkDialTimeout, EnableCompression: true}
	headers := routeWorkspaceHeaders(workspaceID)
	headers.Set("Authorization", "Bearer "+strings.TrimSpace(msg.Token))
	if c.elasticDialContext.Dialer != nil {
		dialer = *c.elasticDialContext.Dialer
	}
	for key, values := range c.elasticDialContext.Headers {
		for _, value := range values {
			headers.Add(key, value)
		}
	}

	conn, resp, err := dialer.DialContext(dialCtx, wsURL, headers)
	if err != nil {
		status := ""
		if resp != nil {
			status = resp.Status
		}
		slog.Warn("AgentTeam: Elastic Request Link dial failed", "channel", c.id, "linkId", msg.LinkID, "status", status, "reason", "dial_failed")
		return
	}
	linkCtx, linkCancel := context.WithCancel(ctx)
	link := &elasticRequestLink{
		linkID:           strings.TrimSpace(msg.LinkID),
		conn:             conn,
		allowedEndpoints: allowedEndpointSet(msg.AllowedEndpoints),
		cancel:           linkCancel,
		done:             make(chan struct{}),
	}
	c.trackElasticRequestLink(link)
	if !c.startConnectionRunTask(linkCtx, func(ctx context.Context) {
		c.runElasticRequestLink(ctx, link, msg)
	}) {
		// The connection run closed after the dial succeeded but before the link
		// runner could publish. Compensate synchronously so neither a carrier nor a
		// link.done waiter escapes the terminal admission boundary.
		link.close()
		c.untrackElasticRequestLink(link)
		close(link.done)
	}
}

func buildElasticRequestLinkURL(serverURL string, connectURL string) (string, error) {
	connectURL = strings.TrimSpace(connectURL)
	if connectURL == "" {
		return "", errors.New("empty connect url")
	}
	if parsed, err := url.Parse(connectURL); err == nil && parsed.IsAbs() {
		if parsed.Query().Has("key") {
			return "", errors.New("elastic request link URL must not carry machine key")
		}
		switch parsed.Scheme {
		case "http":
			parsed.Scheme = "ws"
		case "https":
			parsed.Scheme = "wss"
		}
		if parsed.Scheme != "ws" && parsed.Scheme != "wss" {
			return "", fmt.Errorf("unsupported scheme %q", parsed.Scheme)
		}
		return parsed.String(), nil
	}

	base, err := url.Parse(strings.TrimSpace(serverURL))
	if err != nil {
		return "", err
	}
	switch base.Scheme {
	case "http":
		base.Scheme = "ws"
	case "https":
		base.Scheme = "wss"
	case "ws", "wss":
	default:
		return "", fmt.Errorf("unsupported server scheme %q", base.Scheme)
	}
	ref, err := url.Parse(connectURL)
	if err != nil {
		return "", err
	}
	if ref.Query().Has("key") {
		return "", errors.New("elastic request link URL must not carry machine key")
	}
	if !strings.HasPrefix(ref.Path, "/") {
		ref.Path = "/" + strings.TrimLeft(ref.Path, "/")
	}
	ref.Scheme = ""
	ref.Host = ""
	return base.ResolveReference(ref).String(), nil
}