package oauth

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

	"github.com/lestrrat-go/jwx/v2/jwk"
)

// Client handles atProto OAuth flows (PAR, PKCE, DPoP)
type Client struct {
	clientID    string
	clientJWK   jwk.Key
	redirectURI string
	httpClient  *http.Client
}

// NewClient creates a new OAuth client
func NewClient(clientID string, clientJWK jwk.Key, redirectURI string) *Client {
	return &Client{
		clientID:    clientID,
		clientJWK:   clientJWK,
		redirectURI: redirectURI,
		httpClient: &http.Client{
			Timeout: 30 * time.Second,
		},
	}
}

// AuthServerMetadata represents OAuth 2.0 authorization server metadata (RFC 8414)
type AuthServerMetadata struct {
	Issuer                        string   `json:"issuer"`
	AuthorizationEndpoint         string   `json:"authorization_endpoint"`
	TokenEndpoint                 string   `json:"token_endpoint"`
	PushedAuthReqEndpoint         string   `json:"pushed_authorization_request_endpoint"`
	JWKSURI                       string   `json:"jwks_uri"`
	GrantTypesSupported           []string `json:"grant_types_supported"`
	ResponseTypesSupported        []string `json:"response_types_supported"`
	CodeChallengeMethodsSupported []string `json:"code_challenge_methods_supported"`
	DPoPSigningAlgValuesSupported []string `json:"dpop_signing_alg_values_supported"`
}

// ResolvePDSAuthServer resolves the authorization server for a PDS
// Follows the PDS → Authorization Server discovery flow
func (c *Client) ResolvePDSAuthServer(ctx context.Context, pdsURL string) (string, error) {
	// Fetch PDS metadata from /.well-known/oauth-protected-resource
	metadataURL := strings.TrimSuffix(pdsURL, "/") + "/.well-known/oauth-protected-resource"

	req, err := http.NewRequestWithContext(ctx, "GET", metadataURL, nil)
	if err != nil {
		return "", fmt.Errorf("failed to create request: %w", err)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("failed to fetch PDS metadata: %w", err)
	}
	defer func() { _ = resp.Body.Close() }() //nolint:errcheck

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("PDS returned status %d", resp.StatusCode)
	}

	var metadata struct {
		AuthorizationServers []string `json:"authorization_servers"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&metadata); err != nil {
		return "", fmt.Errorf("failed to decode PDS metadata: %w", err)
	}

	if len(metadata.AuthorizationServers) == 0 {
		return "", fmt.Errorf("no authorization servers found for PDS")
	}

	// Return the first (primary) authorization server
	return metadata.AuthorizationServers[0], nil
}

// FetchAuthServerMetadata fetches OAuth 2.0 authorization server metadata
func (c *Client) FetchAuthServerMetadata(ctx context.Context, issuer string) (*AuthServerMetadata, error) {
	// OAuth 2.0 discovery endpoint
	metadataURL := strings.TrimSuffix(issuer, "/") + "/.well-known/oauth-authorization-server"

	req, err := http.NewRequestWithContext(ctx, "GET", metadataURL, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch auth server metadata: %w", err)
	}
	defer func() { _ = resp.Body.Close() }() //nolint:errcheck

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("auth server returned status %d", resp.StatusCode)
	}

	var metadata AuthServerMetadata
	if err := json.NewDecoder(resp.Body).Decode(&metadata); err != nil {
		return nil, fmt.Errorf("failed to decode auth server metadata: %w", err)
	}

	return &metadata, nil
}

// PARResponse represents the response from a Pushed Authorization Request
type PARResponse struct {
	RequestURI          string `json:"request_uri"`
	ExpiresIn           int    `json:"expires_in"`
	State               string // Generated by client
	PKCEVerifier        string // Generated by client
	DpopAuthserverNonce string // From response header (if provided)
}

// SendPARRequest sends a Pushed Authorization Request (PAR) - RFC 9126
// This pre-registers the authorization request with the server
func (c *Client) SendPARRequest(ctx context.Context, authMeta *AuthServerMetadata, handle, scope string, dpopKey jwk.Key) (*PARResponse, error) {
	// Generate PKCE challenge
	pkce, err := GeneratePKCEChallenge()
	if err != nil {
		return nil, fmt.Errorf("failed to generate PKCE: %w", err)
	}

	// Generate state
	state, err := GenerateState()
	if err != nil {
		return nil, fmt.Errorf("failed to generate state: %w", err)
	}

	// Create form data
	data := url.Values{}
	data.Set("client_id", c.clientID)
	data.Set("redirect_uri", c.redirectURI)
	data.Set("response_type", "code")
	data.Set("scope", scope)
	data.Set("state", state)
	data.Set("code_challenge", pkce.Challenge)
	data.Set("code_challenge_method", pkce.Method)
	data.Set("login_hint", handle) // atProto-specific: suggests which account to use

	// Create DPoP proof for PAR endpoint
	dpopProof, err := CreateDPoPProof(dpopKey, "POST", authMeta.PushedAuthReqEndpoint, "", "")
	if err != nil {
		return nil, fmt.Errorf("failed to create DPoP proof: %w", err)
	}

	// Send PAR request
	req, err := http.NewRequestWithContext(ctx, "POST", authMeta.PushedAuthReqEndpoint, strings.NewReader(data.Encode()))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("DPoP", dpopProof)

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed to send PAR request: %w", err)
	}
	defer func() { _ = resp.Body.Close() }() //nolint:errcheck

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read PAR response: %w", err)
	}

	// Handle DPoP nonce requirement (RFC 9449 Section 8)
	// If server returns use_dpop_nonce error, retry with the nonce
	if resp.StatusCode == http.StatusBadRequest {
		var errorResp struct {
			Error            string `json:"error"`
			ErrorDescription string `json:"error_description"`
		}
		if err := json.Unmarshal(body, &errorResp); err == nil && errorResp.Error == "use_dpop_nonce" {
			// Get nonce from response header
			nonce := resp.Header.Get("DPoP-Nonce")
			if nonce != "" {
				// Retry with nonce
				dpopProof, err = CreateDPoPProof(dpopKey, "POST", authMeta.PushedAuthReqEndpoint, nonce, "")
				if err != nil {
					return nil, fmt.Errorf("failed to create DPoP proof with nonce: %w", err)
				}

				// Re-create request with new DPoP proof
				req, err = http.NewRequestWithContext(ctx, "POST", authMeta.PushedAuthReqEndpoint, strings.NewReader(data.Encode()))
				if err != nil {
					return nil, fmt.Errorf("failed to create retry request: %w", err)
				}
				req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
				req.Header.Set("DPoP", dpopProof)

				// Send retry request
				resp, err = c.httpClient.Do(req)
				if err != nil {
					return nil, fmt.Errorf("failed to send retry PAR request: %w", err)
				}
				defer func() { _ = resp.Body.Close() }() //nolint:errcheck
				body, err = io.ReadAll(resp.Body)
				if err != nil {
					return nil, fmt.Errorf("failed to read retry PAR response: %w", err)
				}
			}
		}
	}

	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated {
		return nil, fmt.Errorf("PAR request failed with status %d", resp.StatusCode)
	}

	var parResp struct {
		RequestURI string `json:"request_uri"`
		ExpiresIn  int    `json:"expires_in"`
	}

	if err := json.Unmarshal(body, &parResp); err != nil {
		return nil, fmt.Errorf("failed to decode PAR response: %w", err)
	}

	// Extract DPoP nonce from response header (if provided)
	dpopNonce := resp.Header.Get("DPoP-Nonce")

	return &PARResponse{
		RequestURI:          parResp.RequestURI,
		ExpiresIn:           parResp.ExpiresIn,
		State:               state,
		PKCEVerifier:        pkce.Verifier,
		DpopAuthserverNonce: dpopNonce,
	}, nil
}

// TokenResponse represents an OAuth token response
type TokenResponse struct {
	AccessToken         string `json:"access_token"`
	TokenType           string `json:"token_type"` // Should be "DPoP"
	ExpiresIn           int    `json:"expires_in"`
	RefreshToken        string `json:"refresh_token"`
	Scope               string `json:"scope"`
	Sub                 string `json:"sub"` // DID of the user
	DpopAuthserverNonce string // From response header
}

// InitialTokenRequest exchanges authorization code for tokens (DPoP-bound)
func (c *Client) InitialTokenRequest(ctx context.Context, code, issuer, pkceVerifier, dpopNonce string, dpopKey jwk.Key) (*TokenResponse, error) {
	// Get auth server metadata for token endpoint
	authMeta, err := c.FetchAuthServerMetadata(ctx, issuer)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch auth server metadata: %w", err)
	}

	// Create form data
	data := url.Values{}
	data.Set("grant_type", "authorization_code")
	data.Set("code", code)
	data.Set("redirect_uri", c.redirectURI)
	data.Set("code_verifier", pkceVerifier)
	data.Set("client_id", c.clientID)

	// Create DPoP proof for token endpoint
	dpopProof, err := CreateDPoPProof(dpopKey, "POST", authMeta.TokenEndpoint, dpopNonce, "")
	if err != nil {
		return nil, fmt.Errorf("failed to create DPoP proof: %w", err)
	}

	// Send token request
	req, err := http.NewRequestWithContext(ctx, "POST", authMeta.TokenEndpoint, strings.NewReader(data.Encode()))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("DPoP", dpopProof)

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed to send token request: %w", err)
	}
	defer func() { _ = resp.Body.Close() }() //nolint:errcheck

	var tokenResp TokenResponse
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("token request failed with status %d", resp.StatusCode)
	}

	if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
		return nil, fmt.Errorf("failed to decode token response: %w", err)
	}

	// Extract updated DPoP nonce
	tokenResp.DpopAuthserverNonce = resp.Header.Get("DPoP-Nonce")

	return &tokenResp, nil
}

// RefreshTokenRequest refreshes an access token using a refresh token
func (c *Client) RefreshTokenRequest(ctx context.Context, refreshToken, issuer, dpopNonce string, dpopKey jwk.Key) (*TokenResponse, error) {
	// Get auth server metadata for token endpoint
	authMeta, err := c.FetchAuthServerMetadata(ctx, issuer)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch auth server metadata: %w", err)
	}

	// Create form data
	data := url.Values{}
	data.Set("grant_type", "refresh_token")
	data.Set("refresh_token", refreshToken)
	data.Set("client_id", c.clientID)

	// Create DPoP proof for token endpoint
	dpopProof, err := CreateDPoPProof(dpopKey, "POST", authMeta.TokenEndpoint, dpopNonce, "")
	if err != nil {
		return nil, fmt.Errorf("failed to create DPoP proof: %w", err)
	}

	// Send refresh request
	req, err := http.NewRequestWithContext(ctx, "POST", authMeta.TokenEndpoint, strings.NewReader(data.Encode()))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("DPoP", dpopProof)

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed to send refresh request: %w", err)
	}
	defer func() { _ = resp.Body.Close() }() //nolint:errcheck

	var tokenResp TokenResponse
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("refresh request failed with status %d", resp.StatusCode)
	}

	if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
		return nil, fmt.Errorf("failed to decode token response: %w", err)
	}

	// Extract updated DPoP nonce
	tokenResp.DpopAuthserverNonce = resp.Header.Get("DPoP-Nonce")

	return &tokenResp, nil
}