package gateway

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

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

// ErrWorkspaceIdentityUnresolved marks a request whose workspace cannot be
// authenticated without guessing.
var ErrWorkspaceIdentityUnresolved = errors.New("workspace identity unresolved")

// ErrWorkspaceIdentityUnauthorized marks a public credential rejected by the
// trusted resolver. It is distinct from the Gateway-to-resolver service token.
var ErrWorkspaceIdentityUnauthorized = errors.New("workspace identity unauthorized")

type WorkspaceIdentityRequest struct {
	Method   string `json:"method"`
	Path     string `json:"path"`
	RawPath  string `json:"rawPath,omitempty"`
	RawQuery string `json:"rawQuery,omitempty"`
	// WorkspaceQualifiedSessionAuthorization carries one Edge-validated session wire
	// value. Keep the private JSON key stable and never log or echo this field.
	WorkspaceQualifiedSessionAuthorization string `json:"authorization,omitempty"`
	Authority                              string `json:"authority,omitempty"`
	WorkspaceID                            string `json:"workspaceId,omitempty"`
}

type RouteIdentity struct {
	WorkspaceID string
}

type WorkspaceIdentityResolver interface {
	ResolveWorkspaceIdentity(context.Context, WorkspaceIdentityRequest) (RouteIdentity, error)
}

type HTTPWorkspaceIdentityResolver struct {
	endpoint *url.URL
	token    string
	client   *http.Client
}

func NewHTTPWorkspaceIdentityResolver(rawURL, token string, client *http.Client) (*HTTPWorkspaceIdentityResolver, error) {
	endpoint, err := parseRequiredUpstream("workspace identity resolver", rawURL)
	if err != nil {
		return nil, err
	}
	if strings.TrimSpace(token) == "" {
		return nil, fmt.Errorf("gateway workspace identity resolver token is required")
	}
	if client == nil {
		client = &http.Client{Timeout: 3 * time.Second}
	}
	return &HTTPWorkspaceIdentityResolver{endpoint: endpoint, token: strings.TrimSpace(token), client: client}, nil
}

func (resolver *HTTPWorkspaceIdentityResolver) ResolveWorkspaceIdentity(ctx context.Context, input WorkspaceIdentityRequest) (RouteIdentity, error) {
	if resolver == nil || resolver.endpoint == nil || resolver.client == nil {
		return RouteIdentity{}, fmt.Errorf("workspace identity resolver is not configured")
	}
	body, err := json.Marshal(input)
	if err != nil {
		return RouteIdentity{}, fmt.Errorf("encode workspace identity request: %w", err)
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, resolver.endpoint.String(), bytes.NewReader(body))
	if err != nil {
		return RouteIdentity{}, fmt.Errorf("build workspace identity request: %w", err)
	}
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set(api.HeaderInternalAdminToken, resolver.token)
	response, err := resolver.client.Do(request)
	if err != nil {
		return RouteIdentity{}, fmt.Errorf("call workspace identity resolver: %w", err)
	}
	defer response.Body.Close()
	if response.StatusCode == http.StatusUnauthorized {
		var failure struct {
			Code string `json:"code"`
		}
		_ = json.NewDecoder(io.LimitReader(response.Body, 4<<10)).Decode(&failure)
		if failure.Code == "workspace_identity_unauthorized" {
			return RouteIdentity{}, ErrWorkspaceIdentityUnauthorized
		}
		return RouteIdentity{}, fmt.Errorf("workspace identity resolver service credential was rejected")
	}
	if response.StatusCode == http.StatusBadRequest || response.StatusCode == http.StatusUnprocessableEntity {
		_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 4<<10))
		return RouteIdentity{}, ErrWorkspaceIdentityUnresolved
	}
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 4<<10))
		return RouteIdentity{}, fmt.Errorf("workspace identity resolver returned status %d", response.StatusCode)
	}
	var output struct {
		WorkspaceID string `json:"workspaceId"`
	}
	if err := json.NewDecoder(io.LimitReader(response.Body, 16<<10)).Decode(&output); err != nil {
		return RouteIdentity{}, fmt.Errorf("decode workspace identity response: %w", err)
	}
	workspaceID := strings.TrimSpace(output.WorkspaceID)
	if workspaceID == "" {
		return RouteIdentity{}, ErrWorkspaceIdentityUnresolved
	}
	return RouteIdentity{WorkspaceID: workspaceID}, nil
}