package auth

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type OAuthConfig struct {
	ClientID    string
	RedirectURI string
	Scopes      []string
}

type OAuthState struct {
	State        string
	CodeVerifier string
	DID          string
	PDS          string
	AccessToken  string
	RefreshToken string
	ExpiresAt    time.Time
	KeyPair      *DPoPKeyPair
}

type OAuthMetadata struct {
	Issuer                string   `json:"issuer"`
	AuthorizationEndpoint string   `json:"authorization_endpoint"`
	TokenEndpoint         string   `json:"token_endpoint"`
	Scopes                []string `json:"scopes_supported"`
}

type TokenResponse struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
	TokenType    string `json:"token_type"`
	ExpiresIn    int    `json:"expires_in"`
}

func (s *OAuthState) Copy() *OAuthState {
	return &OAuthState{
		State:        s.State,
		DID:          s.DID,
		PDS:          s.PDS,
		AccessToken:  s.AccessToken,
		RefreshToken: s.RefreshToken,
		ExpiresAt:    s.ExpiresAt,
		KeyPair:      s.KeyPair,
	}
}

func GenerateRandomToken() (string, error) {
	b := make([]byte, 32)
	_, err := rand.Read(b)
	if err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(b), nil
}

func GenerateCodeChallenge(verifier string) string {
	h := sha256.Sum256([]byte(verifier))
	return base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString(h[:])
}

func FetchOAuthMetadata(pdsURL string) (*OAuthMetadata, error) {
	// first get the resource metadata to find the auth server
	resourceURL := fmt.Sprintf("%s/.well-known/oauth-protected-resource", pdsURL)
	fmt.Printf("Fetching oauth metadata from: %s\n", resourceURL)
	resp, err := http.Get(resourceURL)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	var resource struct {
		AuthorizationServers []string `json:"authorization_servers"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&resource); err != nil {
		return nil, err
	}

	if len(resource.AuthorizationServers) == 0 {
		return nil, fmt.Errorf("no authorization servers found")
	}

	// now fetch the actual oauth metadata from the auth server
	authServer := resource.AuthorizationServers[0]
	metadataURL := fmt.Sprintf("%s/.well-known/oauth-authorization-server", authServer)

	resp, err = http.Get(metadataURL)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	var metadata OAuthMetadata
	if err := json.NewDecoder(resp.Body).Decode(&metadata); err != nil {
		return nil, err
	}

	return &metadata, nil
}

func BuildAuthorizationURL(metadata *OAuthMetadata, config *OAuthConfig, state, codeChallenge, handle string) string {
	params := url.Values{
		"response_type":         {"code"},
		"client_id":             {config.ClientID},
		"redirect_uri":          {config.RedirectURI},
		"state":                 {state},
		"code_challenge":        {codeChallenge},
		"code_challenge_method": {"S256"},
		"scope":                 {"atproto"},
		"login_hint":            {handle},
	}

	return metadata.AuthorizationEndpoint + "?" + params.Encode()
}

func postWithDPoP(url string, data url.Values, keypair *DPoPKeyPair) (*TokenResponse, error) {
	dpopProof, err := keypair.GenerateDPoPProof("POST", url, "")
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest("POST", url, strings.NewReader(data.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("DPoP", dpopProof)

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode == 400 {
		body, _ := io.ReadAll(resp.Body)
		bodyStr := string(body)

		if strings.Contains(bodyStr, "use_dpop_nonce") {
			nonce := resp.Header.Get("DPoP-Nonce")
			if nonce == "" {
				return nil, fmt.Errorf("server requires nonce but didn't provide one")
			}

			dpopProof, err = keypair.GenerateDPoPProof("POST", url, nonce)
			if err != nil {
				return nil, err
			}

			retryReq, err := http.NewRequest("POST", url, strings.NewReader(data.Encode()))
			if err != nil {
				return nil, err
			}
			retryReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
			retryReq.Header.Set("DPoP", dpopProof)

			resp, err = http.DefaultClient.Do(retryReq)
			if err != nil {
				return nil, err
			}
			defer resp.Body.Close()
		} else {
			return nil, fmt.Errorf("request failed: 400 - %s", bodyStr)
		}
	}

	if resp.StatusCode != 200 {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("request failed: %d - %s", resp.StatusCode, string(body))
	}

	var tokens TokenResponse
	if err := json.NewDecoder(resp.Body).Decode(&tokens); err != nil {
		return nil, err
	}

	return &tokens, nil
}

func ExchangeCodeForTokens(pdsURL, code, verifier, callbackURL string, keypair *DPoPKeyPair) (*TokenResponse, error) {
	metadata, err := FetchOAuthMetadata(pdsURL)
	if err != nil {
		return nil, err
	}

	data := url.Values{
		"grant_type":    {"authorization_code"},
		"code":          {code},
		"redirect_uri":  {callbackURL},
		"code_verifier": {verifier},
		"client_id":     {callbackURL},
	}

	return postWithDPoP(metadata.TokenEndpoint, data, keypair)
}

func RefreshAccessToken(pdsURL, refreshToken, callbackURL string, keypair *DPoPKeyPair) (*TokenResponse, error) {
	metadata, err := FetchOAuthMetadata(pdsURL)
	if err != nil {
		return nil, err
	}

	data := url.Values{
		"grant_type":    {"refresh_token"},
		"refresh_token": {refreshToken},
		"client_id":     {callbackURL},
	}

	return postWithDPoP(metadata.TokenEndpoint, data, keypair)
}