package atp

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"github.com/bluesky-social/indigo/atproto/atcrypto"
	"github.com/bluesky-social/indigo/atproto/auth/oauth"
	"github.com/bluesky-social/indigo/atproto/identity"
	"github.com/bluesky-social/indigo/atproto/syntax"
	"github.com/google/go-querystring/query"
	"github.com/pkg/browser"
)

// OAuthConfig configures an [OAuthApp].
type OAuthConfig struct {
	// ClientID is the OAuth client ID (a URL). Empty string or a localhost/127.0.0.1
	// prefix triggers development mode (localhost OAuth with automatic PKCE).
	ClientID string

	// RedirectURI is the fully-qualified OAuth callback URL.
	RedirectURI string

	// Scopes for the OAuth grant. Use [ScopesForCollections] to build these
	// from collection NSIDs. Must include "atproto".
	Scopes []string

	// Store persists OAuth sessions and auth request state. If nil, an
	// in-memory store is used (sessions lost on restart).
	Store oauth.ClientAuthStore

	// AppName is an optional human-readable name included in client metadata.
	AppName string

	// Directory overrides the identity directory used to resolve handles and
	// DIDs during the OAuth flow. If nil, indigo's default directory is used
	// (which queries plc.directory and the public DNS). Set this to point at
	// a test PLC (e.g. for integration tests against an ephemeral PDS).
	Directory identity.Directory
}

// SessionInfo is returned after a successful OAuth callback.
type SessionInfo struct {
	DID       syntax.DID
	SessionID string
	Scopes    []string
}

// OAuthApp manages AT Protocol OAuth flows.
type OAuthApp struct {
	app     *oauth.ClientApp
	appName string
}

// NewOAuthApp creates a new OAuth application from the given config.
func NewOAuthApp(cfg OAuthConfig) (*OAuthApp, error) {
	var config oauth.ClientConfig

	if cfg.ClientID == "" ||
		strings.HasPrefix(cfg.ClientID, "http://localhost") ||
		strings.HasPrefix(cfg.ClientID, "http://127.0.0.1") {
		config = oauth.NewLocalhostConfig(cfg.RedirectURI, cfg.Scopes)
	} else {
		config = oauth.NewPublicConfig(cfg.ClientID, cfg.RedirectURI, cfg.Scopes)
	}

	store := cfg.Store
	if store == nil {
		store = oauth.NewMemStore()
	}

	app := oauth.NewClientApp(&config, store)
	if cfg.Directory != nil {
		app.Dir = cfg.Directory
	}

	return &OAuthApp{
		app:     app,
		appName: cfg.AppName,
	}, nil
}

// StartLogin begins the web OAuth flow. Returns the authorization URL to
// redirect the user to. Use [HandleCallback] when the user returns.
func (a *OAuthApp) StartLogin(ctx context.Context, handle string) (string, error) {
	authURL, err := a.app.StartAuthFlow(ctx, handle)
	if err != nil {
		return "", fmt.Errorf("start auth flow: %w", err)
	}
	return authURL, nil
}

// StartLoginWithScopes is like [StartLogin] but requests the given scopes
// instead of the default [OAuthConfig.Scopes]. Use this for incremental
// scope upgrades: configure OAuthConfig.Scopes with the full superset that
// is declared in client metadata, then call this with a subset for the
// default login and the superset when the user opts in to elevated
// permissions.
//
// The requested scopes must all be a subset of the scopes declared in
// client metadata, or the auth server will reject the request.
// TODO: maybe just make this the default behavior for oauth, seems like a good
// pattern to encourage. (e.g. incrementally add scopes where they are needed)
func (a *OAuthApp) StartLoginWithScopes(ctx context.Context, handle string, scopes []string) (string, error) {
	app := a.app

	atid, err := syntax.ParseAtIdentifier(handle)
	if err != nil {
		return "", fmt.Errorf("not a valid account identifier (%s): %w", handle, err)
	}
	ident, err := app.Dir.Lookup(ctx, atid)
	if err != nil {
		return "", fmt.Errorf("failed to resolve username (%s): %w", handle, err)
	}
	pdsHost := ident.PDSEndpoint()
	if pdsHost == "" {
		return "", fmt.Errorf("identity does not link to an atproto host (PDS)")
	}
	authserverURL, err := app.Resolver.ResolveAuthServerURL(ctx, pdsHost)
	if err != nil {
		return "", fmt.Errorf("resolving auth server: %w", err)
	}
	authserverMeta, err := app.Resolver.ResolveAuthServerMetadata(ctx, authserverURL)
	if err != nil {
		return "", fmt.Errorf("fetching auth server metadata: %w", err)
	}

	info, err := a.sendAuthRequestWithPrompt(ctx, authserverMeta, scopes, "")
	if err != nil {
		return "", fmt.Errorf("auth request failed: %w", err)
	}
	did := ident.DID
	info.AccountDID = &did

	if err := app.Store.SaveAuthRequestInfo(ctx, *info); err != nil {
		return "", fmt.Errorf("saving auth request: %w", err)
	}

	params := url.Values{}
	params.Set("client_id", app.Config.ClientID)
	params.Set("request_uri", info.RequestURI)

	return fmt.Sprintf("%s?%s", authserverMeta.AuthorizationEndpoint, params.Encode()), nil
}

// SessionScopes returns the scopes granted to the given session, or an
// error if the session cannot be loaded. Useful for gating UI on whether a
// session has elevated permissions without resuming a full client.
func (a *OAuthApp) SessionScopes(ctx context.Context, did syntax.DID, sessionID string) ([]string, error) {
	sd, err := a.app.Store.GetSession(ctx, did, sessionID)
	if err != nil {
		return nil, fmt.Errorf("%w: %w", ErrSessionExpired, err)
	}
	return sd.Scopes, nil
}

// HandleCallback processes the OAuth callback after user authorization.
// Pass r.URL.Query() from the callback request.
func (a *OAuthApp) HandleCallback(ctx context.Context, params url.Values) (*SessionInfo, error) {
	sessData, err := a.app.ProcessCallback(ctx, params)
	if err != nil {
		return nil, fmt.Errorf("process callback: %w", err)
	}
	return &SessionInfo{
		DID:       sessData.AccountDID,
		SessionID: sessData.SessionID,
		Scopes:    sessData.Scopes,
	}, nil
}

// LoginCLI runs a complete loopback OAuth flow for CLI applications.
// It opens the user's browser, starts a temporary HTTP server to receive the
// callback, and blocks until authentication completes.
// TODO: should this be part of the library? probably not? (removeds `browser` dep)
func (a *OAuthApp) LoginCLI(ctx context.Context, handle string) (*SessionInfo, error) {
	authURL, err := a.app.StartAuthFlow(ctx, handle)
	if err != nil {
		return nil, fmt.Errorf("start auth flow: %w", err)
	}

	// Parse callback path and port from the configured redirect URI
	redirectURL, err := url.Parse(a.app.Config.CallbackURL)
	if err != nil {
		return nil, fmt.Errorf("parse redirect URI: %w", err)
	}

	if err := browser.OpenURL(authURL); err != nil {
		fmt.Printf("Open this URL in your browser:\n%s\n", authURL)
	}

	type result struct {
		info *SessionInfo
		err  error
	}
	ch := make(chan result, 1)

	mux := http.NewServeMux()
	mux.HandleFunc(redirectURL.Path, func(w http.ResponseWriter, r *http.Request) {
		sessData, err := a.app.ProcessCallback(r.Context(), r.URL.Query())
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprint(w, "Authentication failed")
			ch <- result{err: fmt.Errorf("process callback: %w", err)}
			return
		}

		w.Header().Set("Content-Type", "text/html")
		fmt.Fprint(w, "<html><body><h1>Authenticated!</h1><p>You can close this tab.</p></body></html>")
		ch <- result{info: &SessionInfo{
			DID:       sessData.AccountDID,
			SessionID: sessData.SessionID,
			Scopes:    sessData.Scopes,
		}}
	})

	addr := redirectURL.Host
	server := &http.Server{Addr: addr, Handler: mux}
	go server.ListenAndServe()

	res := <-ch
	server.Close()

	if res.err != nil {
		return nil, res.err
	}
	return res.info, nil
}

// ResumeSession restores an existing OAuth session and returns a [Client]
// ready for PDS CRUD operations.
func (a *OAuthApp) ResumeSession(ctx context.Context, did syntax.DID, sessionID string) (*Client, error) {
	session, err := a.app.ResumeSession(ctx, did, sessionID)
	if err != nil {
		return nil, fmt.Errorf("%w: %w", ErrSessionExpired, err)
	}
	return NewClient(session.APIClient(), did), nil
}

// DeleteSession removes an OAuth session (for logout).
func (a *OAuthApp) DeleteSession(ctx context.Context, did syntax.DID, sessionID string) error {
	return a.app.Store.DeleteSession(ctx, did, sessionID)
}

// Logout revokes the OAuth tokens and deletes the session.
func (a *OAuthApp) Logout(ctx context.Context, did syntax.DID, sessionID string) error {
	return a.app.Logout(ctx, did, sessionID)
}

// ClientMetadata returns the OAuth client metadata document. Serve this
// at your client_id URL and at /.well-known/oauth-client-metadata.
func (a *OAuthApp) ClientMetadata() oauth.ClientMetadata {
	meta := a.app.Config.ClientMetadata()
	if a.appName != "" {
		meta.ClientName = &a.appName
	}
	return meta
}

// StartSignup starts an OAuth flow with prompt=create for account registration.
// pdsURL is the PDS host URL (e.g. "https://arabica.systems").
// Returns the authorization URL to redirect the user to.
func (a *OAuthApp) StartSignup(ctx context.Context, pdsURL string) (string, error) {
	app := a.app

	authserverURL, err := app.Resolver.ResolveAuthServerURL(ctx, pdsURL)
	if err != nil {
		return "", fmt.Errorf("resolving auth server for %s: %w", pdsURL, err)
	}

	authserverMeta, err := app.Resolver.ResolveAuthServerMetadata(ctx, authserverURL)
	if err != nil {
		return "", fmt.Errorf("fetching auth server metadata: %w", err)
	}

	info, err := a.sendAuthRequestWithPrompt(ctx, authserverMeta, app.Config.Scopes, "create")
	if err != nil {
		return "", fmt.Errorf("auth request failed: %w", err)
	}

	if err := app.Store.SaveAuthRequestInfo(ctx, *info); err != nil {
		return "", fmt.Errorf("saving auth request: %w", err)
	}

	params := url.Values{}
	params.Set("client_id", app.Config.ClientID)
	params.Set("request_uri", info.RequestURI)

	redirectURL := fmt.Sprintf("%s?%s", authserverMeta.AuthorizationEndpoint, params.Encode())
	return redirectURL, nil
}

// sendAuthRequestWithPrompt sends a PAR request with an optional prompt parameter.
func (a *OAuthApp) sendAuthRequestWithPrompt(ctx context.Context, authMeta *oauth.AuthServerMetadata, scopes []string, prompt string) (*oauth.AuthRequestData, error) {
	app := a.app
	parURL := authMeta.PushedAuthorizationRequestEndpoint

	state := secureRandomBase64(16)
	pkceVerifier := secureRandomBase64(48)
	codeChallenge := oauth.S256CodeChallenge(pkceVerifier)

	body := oauth.PushedAuthRequest{
		ClientID:            app.Config.ClientID,
		State:               state,
		RedirectURI:         app.Config.CallbackURL,
		Scope:               strings.Join(scopes, " "),
		ResponseType:        "code",
		CodeChallenge:       codeChallenge,
		CodeChallengeMethod: "S256",
	}

	if prompt != "" {
		body.Prompt = &prompt
	}

	if app.Config.IsConfidential() {
		assertionJWT, err := app.Config.NewClientAssertion(authMeta.Issuer)
		if err != nil {
			return nil, err
		}
		body.ClientAssertionType = oauth.ClientAssertionJWTBearer
		body.ClientAssertion = assertionJWT
	}

	vals, err := query.Values(body)
	if err != nil {
		return nil, err
	}
	bodyBytes := []byte(vals.Encode())

	dpopServerNonce := ""
	dpopPrivKey, err := atcrypto.GeneratePrivateKeyP256()
	if err != nil {
		return nil, err
	}

	var resp *http.Response
	for range 2 {
		dpopJWT, err := oauth.NewAuthDPoP("POST", parURL, dpopServerNonce, dpopPrivKey)
		if err != nil {
			return nil, err
		}

		req, err := http.NewRequestWithContext(ctx, "POST", parURL, bytes.NewBuffer(bodyBytes))
		if err != nil {
			return nil, err
		}
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
		req.Header.Set("DPoP", dpopJWT)

		client := app.Client
		if client == nil {
			client = http.DefaultClient
		}
		resp, err = client.Do(req)
		if err != nil {
			return nil, err
		}

		dpopServerNonce = resp.Header.Get("DPoP-Nonce")

		if resp.StatusCode == http.StatusBadRequest && dpopServerNonce != "" {
			var errBody struct {
				Error string `json:"error"`
			}
			bodyData := mustReadBody(resp)
			_ = json.Unmarshal(bodyData, &errBody)
			if errBody.Error == "use_dpop_nonce" {
				continue
			}
			return nil, fmt.Errorf("PAR request failed (HTTP %d): %s", resp.StatusCode, string(bodyData))
		}

		break
	}

	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated {
		bodyData := mustReadBody(resp)
		return nil, fmt.Errorf("PAR request failed (HTTP %d): %s", resp.StatusCode, string(bodyData))
	}

	var parResp oauth.PushedAuthResponse
	if err := json.NewDecoder(resp.Body).Decode(&parResp); err != nil {
		return nil, fmt.Errorf("PAR response decode failed: %w", err)
	}

	info := &oauth.AuthRequestData{
		State:                        state,
		AuthServerURL:                authMeta.Issuer,
		Scopes:                       scopes,
		PKCEVerifier:                 pkceVerifier,
		RequestURI:                   parResp.RequestURI,
		AuthServerTokenEndpoint:      authMeta.TokenEndpoint,
		AuthServerRevocationEndpoint: authMeta.RevocationEndpoint,
		DPoPAuthServerNonce:          dpopServerNonce,
		DPoPPrivateKeyMultibase:      dpopPrivKey.Multibase(),
	}

	return info, nil
}

// secureRandomBase64 generates a cryptographically random base64url-encoded string.
func secureRandomBase64(sizeBytes uint) string {
	b := make([]byte, sizeBytes)
	_, _ = rand.Read(b)
	return base64.RawURLEncoding.EncodeToString(b)
}

// mustReadBody reads and closes the response body, returning the bytes.
func mustReadBody(resp *http.Response) []byte {
	defer resp.Body.Close()
	var buf bytes.Buffer
	buf.ReadFrom(resp.Body)
	return buf.Bytes()
}

func (a *OAuthApp) Store() oauth.ClientAuthStore {
	return a.app.Store
}