package forwarder

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

	"github.com/tidbcloud/AgentTeam/internal/api"
)

// TrustedIdentity is a AgentTeam-verified caller that may be projected to an
// upstream trusted JWT outside of the browser session middleware path.
type TrustedIdentity struct {
	UserID            string
	AgentTeamUserID        string
	Email             string
	Name              string
	PreferredUsername string
	DisplayName       string
	AvatarURL         string
	SessionID         string
}

// TrustedOwnedAgent is an agent ownership projection embedded in an upstream
// trusted JWT for internal requests that operate on agent-owned resources.
type TrustedOwnedAgent struct {
	ID          string `json:"id"`
	Name        string `json:"name,omitempty"`
	DisplayName string `json:"display_name,omitempty"`
}

// JSONRequest describes an internal JSON request that should go through the
// same upstream forwarding/auth path as browser gateway traffic.
type JSONRequest struct {
	Method                    string
	Path                      string
	WorkspaceID               string
	Identity                  TrustedIdentity
	AgentOwnershipWorkspaceID string
	OwnedAgents               []TrustedOwnedAgent
	Body                      any
	Headers                   map[string]string
	// ExpectedStatus optionally requires one exact successful HTTP status.
	// Status 204 is accepted only when it is explicitly expected.
	ExpectedStatus int
	// InternalPurpose is encoded only into ForwardJSON-issued trusted claims.
	// Browser reverse-proxy traffic cannot set or inherit this claim.
	InternalPurpose string
}

// UpstreamHTTPError preserves an upstream non-2xx response while still mapping
// common status codes onto AgentTeam API error sentinels.
type UpstreamHTTPError struct {
	StatusCode int
	Code       string
	Message    string
	rawBody    []byte
	kind       error
}

func (err UpstreamHTTPError) Error() string {
	if strings.TrimSpace(err.Message) != "" {
		return err.Message
	}
	return http.StatusText(err.StatusCode)
}

func (err UpstreamHTTPError) Unwrap() error {
	return err.kind
}

// ForwardJSON sends a JSON request to the configured upstream with the same
// trusted JWT and workspace headers used by the reverse proxy.
func (handler *Handler) ForwardJSON(ctx context.Context, request JSONRequest, response any) error {
	if handler == nil || handler.baseURL == nil {
		return fmt.Errorf("API forwarder is not configured")
	}
	path := strings.TrimSpace(request.Path)
	if path == "" {
		return api.ErrBadRequest
	}
	rawQuery := ""
	if before, after, ok := strings.Cut(path, "?"); ok {
		path = before
		rawQuery = after
	}
	if !strings.HasPrefix(path, "/") {
		path = "/" + path
	}
	method := strings.ToUpper(strings.TrimSpace(request.Method))
	if method == "" {
		method = http.MethodPost
	}
	var body io.Reader
	if request.Body != nil {
		payload, err := json.Marshal(request.Body)
		if err != nil {
			return fmt.Errorf("encode upstream JSON request: %w", err)
		}
		body = bytes.NewReader(payload)
	}
	endpoint := *handler.baseURL
	endpoint.Path = singleJoiningSlash(handler.baseURL.Path, path)
	endpoint.RawQuery = rawQuery
	req, err := http.NewRequestWithContext(ctx, method, endpoint.String(), body)
	if err != nil {
		return fmt.Errorf("build upstream request: %w", err)
	}
	if request.Body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	req.Header.Set("Accept", "application/json")
	for name, value := range request.Headers {
		if strings.TrimSpace(name) == "" || strings.TrimSpace(value) == "" {
			continue
		}
		if strings.EqualFold(strings.TrimSpace(name), "X-AgentTeam-Internal-Purpose") {
			continue
		}
		req.Header.Set(name, value)
	}
	jwt, err := handler.issueTrustedJWTForJSONRequest(request)
	if err != nil {
		return err
	}
	authHeader := "Bearer " + jwt
	req.Header.Set("Authorization", authHeader)
	req.Header.Set(trustedJWTHeader, authHeader)
	if workspaceID := strings.TrimSpace(request.WorkspaceID); workspaceID != "" {
		req.Header.Set(api.HeaderWorkspaceID, workspaceID)
	} else {
		req.Header.Del(api.HeaderWorkspaceID)
	}

	client := handler.httpClient
	if client == nil {
		client = &http.Client{Timeout: defaultProxyTimeout}
	}
	res, err := client.Do(req)
	if err != nil {
		return fmt.Errorf("forward upstream request: %w", err)
	}
	defer res.Body.Close()
	if res.StatusCode < http.StatusOK || res.StatusCode >= http.StatusMultipleChoices {
		return upstreamHTTPError(res)
	}
	if request.ExpectedStatus != 0 && res.StatusCode != request.ExpectedStatus {
		return fmt.Errorf("unexpected upstream HTTP status: expected %d, received %d", request.ExpectedStatus, res.StatusCode)
	}
	if res.StatusCode == http.StatusNoContent {
		if request.ExpectedStatus != http.StatusNoContent {
			return fmt.Errorf("unexpected upstream HTTP status: received 204 without explicit expectation")
		}
		if response != nil {
			return fmt.Errorf("expected status 204 requires a nil response target")
		}
		return nil
	}
	target := response
	if target == nil {
		target = new(any)
	}
	decoder := json.NewDecoder(res.Body)
	if err := decoder.Decode(target); err != nil {
		return fmt.Errorf("decode upstream JSON response: %w", err)
	}
	var trailing any
	if err := decoder.Decode(&trailing); err != io.EOF {
		if err == nil {
			return fmt.Errorf("decode upstream JSON response: multiple JSON values")
		}
		return fmt.Errorf("decode upstream JSON response: %w", err)
	}
	return nil
}

func (handler *Handler) issueTrustedJWTForJSONRequest(request JSONRequest) (string, error) {
	claims, err := handler.trustedClaimsFromIdentity(request.Identity)
	if err != nil {
		return "", err
	}
	claims.InternalPurpose = strings.TrimSpace(request.InternalPurpose)
	workspaceID := strings.TrimSpace(request.AgentOwnershipWorkspaceID)
	if workspaceID != "" {
		claims.AgentOwnershipWorkspaceID = workspaceID
	}
	if len(request.OwnedAgents) > 0 {
		claims.OwnedAgents = make([]TrustedOwnedAgent, 0, len(request.OwnedAgents))
		claims.OwnedAgentIDs = make([]string, 0, len(request.OwnedAgents))
		seen := make(map[string]struct{}, len(request.OwnedAgents))
		for _, agent := range request.OwnedAgents {
			id := strings.TrimSpace(agent.ID)
			if id == "" {
				continue
			}
			if _, ok := seen[id]; ok {
				continue
			}
			seen[id] = struct{}{}
			claims.OwnedAgentIDs = append(claims.OwnedAgentIDs, id)
			claims.OwnedAgents = append(claims.OwnedAgents, TrustedOwnedAgent{
				ID:          id,
				Name:        strings.TrimSpace(agent.Name),
				DisplayName: strings.TrimSpace(agent.DisplayName),
			})
		}
	}
	return encodeUnverifiedJWT(claims)
}

func (handler *Handler) issueTrustedJWTForIdentity(identity TrustedIdentity) (string, error) {
	claims, err := handler.trustedClaimsFromIdentity(identity)
	if err != nil {
		return "", err
	}
	return encodeUnverifiedJWT(claims)
}

func (handler *Handler) trustedClaimsFromIdentity(identity TrustedIdentity) (trustedClaims, error) {
	now := handler.now
	if now == nil {
		now = time.Now
	}
	issuedAt := now().UTC()
	subject := firstNonEmpty(identity.UserID, identity.AgentTeamUserID)
	if subject == "" {
		return trustedClaims{}, fmt.Errorf("authenticated user subject is empty")
	}
	agentteamUserID := firstNonEmpty(identity.AgentTeamUserID, identity.UserID)
	preferredUsername := firstNonEmpty(
		identity.PreferredUsername,
		normalizePreferredUsernameCandidate(identity.Name),
		normalizePreferredUsernameCandidate(emailLocalPart(identity.Email)),
		normalizePreferredUsernameCandidate(identity.DisplayName),
		normalizePreferredUsernameCandidate(agentteamUserID),
	)
	return trustedClaims{
		Issuer:            "agentteam",
		Audience:          handler.trustedJWTAudience,
		Subject:           subject,
		AgentTeamUserID:        agentteamUserID,
		Email:             identity.Email,
		Name:              identity.Name,
		PreferredUsername: preferredUsername,
		DisplayName:       identity.DisplayName,
		AvatarURL:         identity.AvatarURL,
		SessionID:         identity.SessionID,
		IssuedAt:          issuedAt.Unix(),
		ExpiresAt:         issuedAt.Add(5 * time.Minute).Unix(),
	}, nil
}

func upstreamHTTPError(response *http.Response) error {
	var payload struct {
		Code    string `json:"code"`
		Message string `json:"message"`
		Error   string `json:"error"`
	}
	body, _ := io.ReadAll(io.LimitReader(response.Body, 1<<20))
	if len(body) > 0 {
		_ = json.Unmarshal(body, &payload)
	}
	return UpstreamHTTPError{
		StatusCode: response.StatusCode,
		Code:       firstNonEmpty(payload.Code, payload.Error),
		Message:    firstNonEmpty(payload.Message, payload.Error, payload.Code, http.StatusText(response.StatusCode)),
		rawBody:    append([]byte(nil), body...),
		kind:       upstreamHTTPErrorKind(response.StatusCode),
	}
}

func upstreamHTTPErrorKind(status int) error {
	switch status {
	case http.StatusUnauthorized:
		return api.ErrUnauthorized
	case http.StatusForbidden:
		return api.ErrForbidden
	case http.StatusNotFound:
		return api.ErrNotFound
	case http.StatusBadRequest, http.StatusUnprocessableEntity:
		return api.ErrBadRequest
	case http.StatusConflict:
		return api.ErrConflict
	default:
		return nil
	}
}