package cognitoauth

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

	"github.com/labstack/echo/v4"
)

// initiateLoginWithPKCE generates PKCE code challenge and redirects to Cognito
// This function starts the OAuth authorization flow with PKCE
func initiateLoginWithPKCE(c echo.Context, config *CognitoConfig) error {
	// Generate random state parameter to prevent CSRF
	state, err := generateRandomString(32)
	if err != nil {
		config.Logger.Error("Failed to generate state parameter", "error", err)
		return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Authentication initialization failed"})
	}

	// Generate code verifier (random string between 43-128 chars)
	codeVerifier, err := generateRandomString(64)
	if err != nil {
		config.Logger.Error("Failed to generate code verifier", "error", err)
		return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Authentication initialization failed"})
	}

	// Create code challenge from verifier (SHA256 + Base64URL without padding)
	codeChallenge := createCodeChallenge(codeVerifier)

	// Store in session for later verification (with expiration time)
	storePKCESession(state, codeVerifier, config.Logger)

	// Build authorization URL with PKCE parameters
	params := url.Values{}
	params.Set("client_id", config.ClientID)
	params.Set("response_type", "code")
	params.Set("redirect_uri", config.RedirectURI)
	params.Set("scope", "openid email profile")
	params.Set("state", state)
	params.Set("code_challenge", codeChallenge)
	params.Set("code_challenge_method", "S256")

	authURL := fmt.Sprintf("%s?%s", config.AuthURL, params.Encode())

	if os.Getenv("DEBUG") == "true" {
		config.Logger.Debug("Initiating login with PKCE",
			"code_verifier", codeVerifier,
			"code_challenge", codeChallenge,
			"state", state,
			"redirect_uri", config.RedirectURI)
	}

	// Redirect user to Cognito login page
	return c.Redirect(http.StatusFound, authURL)
}

// exchangeCodeForTokensWithPKCE exchanges the authorization code for tokens using PKCE
// Called during the OAuth callback flow to get tokens from the authorization code
// Makes an HTTP request to Cognito's token endpoint to perform this exchange
func exchangeCodeForTokensWithPKCE(authCode, codeVerifier string, config CognitoConfig) (*TokenResponse, error) {
	data := url.Values{}
	data.Set("grant_type", "authorization_code")
	data.Set("client_id", config.ClientID)
	data.Set("code", authCode)
	data.Set("redirect_uri", config.RedirectURI)
	data.Set("code_verifier", codeVerifier) // Include code verifier for PKCE

	config.Logger.Debug("Token request details",
		"url", config.TokenURL,
		"client_id", config.ClientID,
		"redirect_uri", config.RedirectURI)

	req, err := http.NewRequest("POST", config.TokenURL, strings.NewReader(data.Encode()))
	if err != nil {
		return nil, fmt.Errorf("failed to create token request: %w", err)
	}

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

	// Add Authorization header if client secret is provided
	if config.ClientSecret != "" {
		req.SetBasicAuth(config.ClientID, config.ClientSecret)
		config.Logger.Debug("Using Basic Auth authentication")
	} else {
		config.Logger.Debug("Using public client authentication (no secret)")
	}

	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("token request failed: %w", err)
	}
	defer resp.Body.Close()

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

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("token endpoint returned status %d: %s", resp.StatusCode, string(body))
	}

	var tokenResponse TokenResponse
	if err := json.Unmarshal(body, &tokenResponse); err != nil {
		return nil, fmt.Errorf("failed to parse token response: %w", err)
	}

	return &tokenResponse, nil
}

// handleOAuthCallback handles the OAuth 2.0 authorization code flow callback with PKCE
// Called when Cognito redirects back to our application with an authorization code
// Exchanges the code for tokens, verifies them, and redirects to home page with token in cookie
func handleOAuthCallback(c echo.Context, config *CognitoConfig) error {
	// Extract the authorization code and state
	code := c.QueryParam("code")
	if code == "" {
		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Missing code parameter"})
	}

	state := c.QueryParam("state")
	if state == "" {
		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Missing state parameter"})
	}

	// Check for OAuth errors
	errorMsg := c.QueryParam("error")
	if errorMsg != "" {
		errorDesc := c.QueryParam("error_description")
		return c.JSON(http.StatusBadRequest, map[string]string{
			"error":             errorMsg,
			"error_description": errorDesc,
		})
	}

	// Retrieve stored code verifier
	codeVerifier, err := getCodeVerifier(state, config.Logger)
	if err != nil {
		config.Logger.Error("Failed to retrieve code verifier", "error", err, "state", state)
		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Invalid or expired session"})
	}

	// Exchange the code for tokens using PKCE
	tokens, err := exchangeCodeForTokensWithPKCE(code, codeVerifier, *config)
	if err != nil {
		config.Logger.Error("Failed to exchange code for tokens", "error", err)
		return c.JSON(http.StatusUnauthorized, map[string]string{"error": err.Error()})
	}

	// Get the JWKS for token verification
	keySet, err := GetJWKS(config.JwksURL, config.Logger)
	if err != nil {
		config.Logger.Error("Failed to fetch JWKS", "error", err)
		return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Failed to verify token"})
	}

	// Verify ID token
	idTokenClaims, err := verifyToken(tokens.IDToken, keySet, *config)
	if err != nil {
		config.Logger.Error("Failed to verify ID token", "error", err)
		return c.JSON(http.StatusUnauthorized, map[string]string{"error": "Invalid token"})
	}

	// Extract user groups
	userGroups, err := GetUserGroups(idTokenClaims)
	if err != nil {
		config.Logger.Warn("Failed to extract user groups", "error", err)
		userGroups = []string{} // Initialize as empty array
	}

	// Debug logging for tokens
	if os.Getenv("DEBUG") == "true" {
		// Log decoded token for debugging
		config.Logger.Info("ID Token Claims", "claims", idTokenClaims)
		config.Logger.Info("Raw ID Token", "token", tokens.IDToken)
		config.Logger.Info("Raw Access Token", "token", tokens.AccessToken)
	}

	// Check authorization
	authorized, requiredGroups := checkPermissions(config.CallbackPath, userGroups, config.RoutePermissions, config.Logger)
	if !authorized {
		config.Logger.Warn("OAuth callback - access denied",
			"user", idTokenClaims["cognito:username"],
			"groups", userGroups,
			"required", requiredGroups)

		return c.JSON(http.StatusForbidden, map[string]interface{}{
			"error":           "Insufficient permissions",
			"message":         "User authenticated successfully but doesn't have the required group membership",
			"authenticated":   true,
			"username":        idTokenClaims["cognito:username"],
			"email":           idTokenClaims["email"],
			"groups":          userGroups,
			"required_groups": requiredGroups,
		})
	}

	// User is authenticated and authorized

	// Calculate token expiration time based on token's expires_in value
	expiresAt := time.Now().Add(time.Duration(tokens.ExpiresIn) * time.Second)

	// Set a cookie with the access token
	tokenCookie := new(http.Cookie)
	tokenCookie.Name = "auth_token"
	tokenCookie.Value = tokens.AccessToken
	tokenCookie.Path = "/"
	tokenCookie.Expires = expiresAt
	tokenCookie.HttpOnly = true // Not accessible via JavaScript
	// In production, set Secure to true
	// tokenCookie.Secure = true
	c.SetCookie(tokenCookie)

	// Extract username for display purposes (optional)
	username, _ := idTokenClaims["cognito:username"].(string)

	config.Logger.Info("User authenticated successfully", "username", username, "groups", userGroups)

	// Redirect to home page
	return c.Redirect(http.StatusFound, config.HomePath)
}

// LogoutHandler handles the logout process by clearing the auth cookie
func LogoutHandler(c echo.Context, config *CognitoConfig) error {
	// Clear the token cookie
	cookie := new(http.Cookie)
	cookie.Name = "auth_token"
	cookie.Value = ""
	cookie.Path = "/"
	cookie.Expires = time.Now().Add(-1 * time.Hour) // Set expiry in the past
	cookie.HttpOnly = true
	c.SetCookie(cookie)

	// Redirect to home page
	return c.Redirect(http.StatusFound, config.HomePath)
}