package wire

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

const daemonTokenRefreshProofHeader = "X-AgentTeam-Daemon-Connection-Token"

type agentTokenRefreshState struct {
	done  chan struct{}
	token string
	err   error
}

// refreshAgentToken resolves a replacement agent token whenever a held token
// fails with 401. It returns a cached newer token first: staggered holders of
// superseded snapshots must reuse the newest issued token, because each
// additional issuance invalidates the token another holder just obtained.
// Only when the cache holds no newer token (concurrent refreshes share one
// in-flight issuance) does it re-issue from the machine-authenticated server
// endpoint. When the cached token itself is also rotated out, the caller's
// bounded second refresh round falls through to a real re-issue.
func (c *Protocol) refreshAgentToken(ctx context.Context, serverAgentID string, staleToken string) (string, error) {
	serverAgentID = strings.TrimSpace(serverAgentID)
	staleToken = strings.TrimSpace(staleToken)
	if serverAgentID == "" {
		return "", fmt.Errorf("agentteam refresh agent auth token: agent id is required")
	}

	for {
		c.mu.Lock()
		if st, ok := c.agents[serverAgentID]; ok {
			if token := strings.TrimSpace(st.Token); token != "" && token != staleToken {
				c.mu.Unlock()
				return token, nil
			}
		}
		if c.agentTokenRefresh == nil {
			c.agentTokenRefresh = make(map[string]*agentTokenRefreshState)
		}
		if inFlight := c.agentTokenRefresh[serverAgentID]; inFlight != nil {
			done := inFlight.done
			c.mu.Unlock()
			select {
			case <-done:
				if inFlight.err != nil {
					return "", inFlight.err
				}
				if token := strings.TrimSpace(inFlight.token); token != "" {
					return token, nil
				}
				// Defensive fallback: retry if an in-flight refresh completed
				// without an error but did not produce a usable token.
				continue
			case <-ctx.Done():
				return "", ctx.Err()
			}
		}

		state := &agentTokenRefreshState{done: make(chan struct{})}
		c.agentTokenRefresh[serverAgentID] = state
		c.mu.Unlock()

		token, err := c.issueAgentAuthToken(ctx, serverAgentID)
		token = strings.TrimSpace(token)
		state.token = token
		state.err = err

		c.mu.Lock()
		if err == nil {
			agentState := c.getOrCreateAgentState(serverAgentID)
			receiptSequence := agentState.nextReceiptSequence()
			agentState.Token = token
			agentState.tokenReceiptSequence = receiptSequence
		}
		if c.agentTokenRefresh[serverAgentID] == state {
			delete(c.agentTokenRefresh, serverAgentID)
		}
		close(state.done)
		c.mu.Unlock()

		if err != nil {
			return "", err
		}
		return token, nil
	}
}

func (c *Protocol) issueAgentAuthToken(ctx context.Context, serverAgentID string) (string, error) {
	httpClient := c.httpClient
	if httpClient == nil {
		httpClient = http.DefaultClient
	}
	serverURL, apiKey, workspaceID, _, _, _ := c.identitySnapshot()
	endpoint := strings.TrimRight(serverURL, "/") + daemonAgentAuthTokenPath(workspaceID, serverAgentID)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, http.NoBody)
	if err != nil {
		return "", fmt.Errorf("agentteam refresh agent auth token: %w", err)
	}
	setRequestHeaders(req, requestHeaders{
		AuthToken:        strings.TrimSpace(apiKey),
		RouteWorkspaceID: workspaceID,
	})
	if proof := c.currentTokenRefreshProof(); proof != "" {
		req.Header.Set(daemonTokenRefreshProofHeader, proof)
	}

	resp, err := httpClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("agentteam refresh agent auth token request: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 400 {
		return "", fmt.Errorf("agentteam refresh agent auth token failed: HTTP %d: %s", resp.StatusCode, readAgentTeamErrorBody(resp.Body))
	}

	var payload struct {
		AuthToken string `json:"authToken"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {
		return "", fmt.Errorf("agentteam refresh agent auth token decode: %w", err)
	}
	token := strings.TrimSpace(payload.AuthToken)
	if token == "" {
		return "", fmt.Errorf("agentteam refresh agent auth token: empty auth token")
	}
	return token, nil
}

func readAgentTeamErrorBody(body io.Reader) string {
	if body == nil {
		return ""
	}
	data, err := io.ReadAll(io.LimitReader(body, 64*1024))
	if err != nil {
		return ""
	}
	var payload struct {
		Error   string `json:"error"`
		Message string `json:"message"`
	}
	if err := json.Unmarshal(data, &payload); err == nil {
		if text := strings.TrimSpace(payload.Error); text != "" {
			return text
		}
		if text := strings.TrimSpace(payload.Message); text != "" {
			return text
		}
	}
	return strings.TrimSpace(string(data))
}