package oauth

import (
	"context"
	"crypto/rand"
	"embed"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"log/slog"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/SigNoz/signoz-mcp-server/internal/client"
	"github.com/SigNoz/signoz-mcp-server/internal/config"
	"github.com/SigNoz/signoz-mcp-server/pkg/analytics"
	logpkg "github.com/SigNoz/signoz-mcp-server/pkg/log"
	otelpkg "github.com/SigNoz/signoz-mcp-server/pkg/otel"
	"github.com/SigNoz/signoz-mcp-server/pkg/util"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/metric"
)

const csrfCookieName = "signoz_mcp_oauth_csrf"

// AnalyticsEmitter publishes OAuth flow events. apiKey and signozURL are
// the validated tenant credentials the emitter needs to resolve identity
// (orgId) for the event. Nil disables OAuth analytics.
type AnalyticsEmitter func(ctx context.Context, event, apiKey, signozURL string, props map[string]any)

//go:embed static/authorize.html
var authorizeTemplateFS embed.FS

var authorizePageTemplate = template.Must(template.ParseFS(authorizeTemplateFS, "static/authorize.html"))

type Handler struct {
	logger            *slog.Logger
	config            *config.Config
	tokenSecret       []byte
	authorizeTemplate *template.Template
	emitEvent         AnalyticsEmitter
	meters            *otelpkg.Meters
}

type registerClientRequest struct {
	ClientName   string   `json:"client_name"`
	RedirectURIs []string `json:"redirect_uris"`
}

type registerClientResponse struct {
	ClientID                string   `json:"client_id"`
	ClientIDIssuedAt        int64    `json:"client_id_issued_at"`
	ClientName              string   `json:"client_name"`
	RedirectURIs            []string `json:"redirect_uris"`
	GrantTypes              []string `json:"grant_types"`
	ResponseTypes           []string `json:"response_types"`
	TokenEndpointAuthMethod string   `json:"token_endpoint_auth_method"`
}

type authorizeTemplateData struct {
	ClientID            string
	ClientName          string
	RedirectURI         string
	AuthorizePath       string
	State               string
	CodeChallenge       string
	CodeChallengeMethod string
	Scope               string
	CSRFToken           string
	SignozURL           string
	ErrorMessage        string
	ErrorCode           string
}

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

func NewHandler(logger *slog.Logger, cfg *config.Config, emitEvent AnalyticsEmitter, meters *otelpkg.Meters) *Handler {
	return &Handler{
		logger:            logger,
		config:            cfg,
		tokenSecret:       []byte(cfg.OAuthTokenSecret),
		authorizeTemplate: authorizePageTemplate,
		emitEvent:         emitEvent,
		meters:            meters,
	}
}

func (h *Handler) emit(ctx context.Context, event, apiKey, signozURL string, props map[string]any) {
	if h.emitEvent == nil {
		return
	}
	h.emitEvent(ctx, event, apiKey, signozURL, props)
}

func (h *Handler) HandleProtectedResourceMetadata(w http.ResponseWriter, r *http.Request) {
	h.writeJSON(w, http.StatusOK, map[string]any{
		"resource":                 h.issuerURL() + "/mcp",
		"authorization_servers":    []string{h.issuerURL()},
		"bearer_methods_supported": []string{"header"},
	})
}

func (h *Handler) HandleAuthorizationServerMetadata(w http.ResponseWriter, r *http.Request) {
	h.writeJSON(w, http.StatusOK, map[string]any{
		"issuer":                                h.issuerURL(),
		"authorization_endpoint":                h.issuerURL() + "/oauth/authorize",
		"token_endpoint":                        h.issuerURL() + "/oauth/token",
		"registration_endpoint":                 h.issuerURL() + "/oauth/register",
		"response_types_supported":              []string{"code"},
		"grant_types_supported":                 []string{"authorization_code", "refresh_token"},
		"code_challenge_methods_supported":      []string{"S256"},
		"token_endpoint_auth_methods_supported": []string{"none"},
	})
}

func (h *Handler) HandleRegisterClient(w http.ResponseWriter, r *http.Request) {
	var req registerClientRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_client_metadata", "request body must be valid JSON")
		return
	}

	req.ClientName = strings.TrimSpace(req.ClientName)
	if req.ClientName == "" {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_client_metadata", "client_name is required")
		return
	}
	if len(req.RedirectURIs) == 0 {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_redirect_uri", "at least one redirect URI is required")
		return
	}

	for _, redirectURI := range req.RedirectURIs {
		if err := validateRedirectURI(redirectURI); err != nil {
			h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_redirect_uri", err.Error())
			return
		}
	}

	createdAt := time.Now().UTC()
	clientID, err := EncryptClientID(req.RedirectURIs, req.ClientName, createdAt, h.tokenSecret)
	if err != nil {
		h.writeOAuthError(r, w, http.StatusInternalServerError, "server_error", "failed to register client")
		return
	}

	h.writeJSON(w, http.StatusCreated, registerClientResponse{
		ClientID:                clientID,
		ClientIDIssuedAt:        createdAt.Unix(),
		ClientName:              req.ClientName,
		RedirectURIs:            req.RedirectURIs,
		GrantTypes:              []string{"authorization_code", "refresh_token"},
		ResponseTypes:           []string{"code"},
		TokenEndpointAuthMethod: "none",
	})
}

func (h *Handler) HandleAuthorizePage(w http.ResponseWriter, r *http.Request) {
	params, err := h.validateAuthorizeRequest(r.URL.Query())
	if err != nil {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_request", err.Error())
		return
	}

	h.renderAuthorizePage(w, r, http.StatusOK, params)
}

func (h *Handler) HandleAuthorizeSubmit(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_request", "form data is required")
		return
	}

	if !h.validateCSRF(r) {
		h.writeOAuthError(r, w, http.StatusForbidden, "access_denied", "invalid CSRF token")
		return
	}

	params, err := h.validateAuthorizeRequest(r.PostForm)
	if err != nil {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_request", err.Error())
		return
	}

	apiKey := strings.TrimSpace(r.FormValue("api_key"))
	signozURL := strings.TrimSpace(r.FormValue("signoz_url"))
	if apiKey == "" {
		h.renderAuthorizePage(w, r, http.StatusBadRequest, authorizeTemplateData{
			ClientID:            params.ClientID,
			ClientName:          params.ClientName,
			RedirectURI:         params.RedirectURI,
			State:               params.State,
			CodeChallenge:       params.CodeChallenge,
			CodeChallengeMethod: params.CodeChallengeMethod,
			Scope:               params.Scope,
			SignozURL:           signozURL,
			ErrorMessage:        "Enter your SigNoz API key to continue.",
			ErrorCode:           "invalid_request",
		})
		return
	}
	normalizedURL, err := util.NormalizeSigNozURLFormInput(signozURL)
	if err != nil {
		h.renderAuthorizePage(w, r, http.StatusBadRequest, authorizeTemplateData{
			ClientID:            params.ClientID,
			ClientName:          params.ClientName,
			RedirectURI:         params.RedirectURI,
			State:               params.State,
			CodeChallenge:       params.CodeChallenge,
			CodeChallengeMethod: params.CodeChallengeMethod,
			Scope:               params.Scope,
			SignozURL:           signozURL,
			ErrorMessage:        "Enter a valid SigNoz URL, for example your-instance.signoz.cloud.",
			ErrorCode:           "invalid_request",
		})
		return
	}
	if err := h.validateSigNozCredentials(r.Context(), normalizedURL, apiKey); err != nil {
		switch {
		case errors.Is(err, client.ErrUnauthorized):
			h.renderAuthorizePage(w, r, http.StatusUnauthorized, authorizeTemplateData{
				ClientID:            params.ClientID,
				ClientName:          params.ClientName,
				RedirectURI:         params.RedirectURI,
				State:               params.State,
				CodeChallenge:       params.CodeChallenge,
				CodeChallengeMethod: params.CodeChallengeMethod,
				Scope:               params.Scope,
				SignozURL:           normalizedURL,
				ErrorMessage:        "We couldn't sign in to that SigNoz instance. Check the URL and API key, then try again.",
				ErrorCode:           "access_denied",
			})
		default:
			h.renderAuthorizePage(w, r, http.StatusBadGateway, authorizeTemplateData{
				ClientID:            params.ClientID,
				ClientName:          params.ClientName,
				RedirectURI:         params.RedirectURI,
				State:               params.State,
				CodeChallenge:       params.CodeChallenge,
				CodeChallengeMethod: params.CodeChallengeMethod,
				Scope:               params.Scope,
				SignozURL:           normalizedURL,
				ErrorMessage:        "We couldn't reach that SigNoz instance. Check the URL and try again.",
				ErrorCode:           "temporarily_unavailable",
			})
		}
		return
	}

	code, err := EncryptAuthorizationCode(
		apiKey,
		normalizedURL,
		params.ClientID,
		params.RedirectURI,
		params.CodeChallenge,
		params.CodeChallengeMethod,
		time.Now().UTC().Add(h.config.AuthCodeTTL),
		h.tokenSecret,
	)
	if err != nil {
		h.writeOAuthError(r, w, http.StatusInternalServerError, "server_error", "failed to generate authorization code")
		return
	}

	redirectURL, err := addAuthorizeResponse(params.RedirectURI, code, params.State)
	if err != nil {
		h.writeOAuthError(r, w, http.StatusInternalServerError, "server_error", "failed to build redirect URL")
		return
	}

	http.SetCookie(w, &http.Cookie{
		Name:     csrfCookieName,
		Value:    "",
		Path:     h.authorizePath(),
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
		Secure:   h.isSecure(),
		MaxAge:   -1,
	})

	h.emit(r.Context(), analytics.EventOAuthAuthorizationSubmitted, apiKey, normalizedURL, map[string]any{
		analytics.AttrTenantURL: normalizedURL,
	})

	http.Redirect(w, r, redirectURL, http.StatusFound)
}

func (h *Handler) validateSigNozCredentials(ctx context.Context, signozURL, apiKey string) error {
	// Only forward custom headers when the user-supplied URL matches the
	// configured SIGNOZ_URL to prevent leaking proxy-auth credentials to
	// attacker-controlled hosts.
	var headers map[string]string
	configNormalized, _ := util.NormalizeSigNozURL(h.config.URL)
	if strings.EqualFold(signozURL, configNormalized) {
		headers = h.config.CustomHeaders
	}
	signozClient := client.NewClient(h.logger, signozURL, apiKey, "SIGNOZ-API-KEY", headers)
	return signozClient.ValidateCredentials(ctx)
}

func (h *Handler) renderAuthorizePage(w http.ResponseWriter, r *http.Request, status int, data authorizeTemplateData) {
	ctx := context.Background()
	if r != nil {
		ctx = r.Context()
	}

	csrfToken, err := randomURLSafeString(32)
	if err != nil {
		h.writeOAuthError(r, w, http.StatusInternalServerError, "server_error", "failed to generate CSRF token")
		return
	}

	http.SetCookie(w, &http.Cookie{
		Name:     csrfCookieName,
		Value:    csrfToken,
		Path:     h.authorizePath(),
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
		Secure:   h.isSecure(),
		MaxAge:   int(h.config.AuthCodeTTL.Seconds()),
	})

	data.AuthorizePath = h.authorizePath()
	data.CSRFToken = csrfToken

	if status >= http.StatusBadRequest && data.ErrorCode != "" {
		h.recordOAuthFailure(ctx, r, status, data.ErrorCode, data.ErrorMessage)
	}

	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.WriteHeader(status)
	if err := h.authorizeTemplate.Execute(w, data); err != nil {
		h.logger.ErrorContext(ctx, "failed to render authorization page", logpkg.ErrAttr(err))
	}
}

func (h *Handler) HandleToken(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_request", "form data is required")
		return
	}

	switch r.FormValue("grant_type") {
	case "authorization_code":
		h.handleAuthorizationCodeGrant(w, r)
	case "refresh_token":
		h.handleRefreshTokenGrant(w, r)
	default:
		h.writeOAuthError(r, w, http.StatusBadRequest, "unsupported_grant_type", "grant_type must be authorization_code or refresh_token")
	}
}

func (h *Handler) handleAuthorizationCodeGrant(w http.ResponseWriter, r *http.Request) {
	clientID := strings.TrimSpace(r.FormValue("client_id"))
	code := strings.TrimSpace(r.FormValue("code"))
	redirectURI := strings.TrimSpace(r.FormValue("redirect_uri"))
	codeVerifier := strings.TrimSpace(r.FormValue("code_verifier"))

	if clientID == "" || code == "" || redirectURI == "" || codeVerifier == "" {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_request", "client_id, code, redirect_uri, and code_verifier are required")
		return
	}

	apiKey, signozURL, authClientID, authRedirectURI, codeChallenge, codeChallengeMethod, _, err := DecryptAuthorizationCode(code, h.tokenSecret)
	if err != nil {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_grant", "authorization code is invalid or expired")
		return
	}
	// Seed tenant_url on ctx so downstream OAuth failure metrics and logs
	// carry the tenant dimension for post-decrypt failure modes.
	r = r.WithContext(util.SetSigNozURL(r.Context(), signozURL))
	if authClientID != clientID || authRedirectURI != redirectURI {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_grant", "authorization code does not match the client or redirect URI")
		return
	}
	if !ValidatePKCE(codeVerifier, codeChallenge, codeChallengeMethod) {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_grant", "PKCE validation failed")
		return
	}

	h.issueTokenPair(w, r, clientID, apiKey, signozURL, "authorization_code")
}

func (h *Handler) handleRefreshTokenGrant(w http.ResponseWriter, r *http.Request) {
	refreshTokenValue := strings.TrimSpace(r.FormValue("refresh_token"))
	clientID := strings.TrimSpace(r.FormValue("client_id"))
	if refreshTokenValue == "" {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_request", "refresh_token is required")
		return
	}

	apiKey, signozURL, refreshClientID, _, err := DecryptRefreshToken(refreshTokenValue, h.tokenSecret)
	if err != nil {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_grant", "refresh token is invalid or expired")
		return
	}
	r = r.WithContext(util.SetSigNozURL(r.Context(), signozURL))
	if clientID != "" && refreshClientID != clientID {
		h.writeOAuthError(r, w, http.StatusBadRequest, "invalid_grant", "refresh token does not belong to this client")
		return
	}

	h.issueTokenPair(w, r, refreshClientID, apiKey, signozURL, "refresh_token")
}

func (h *Handler) issueTokenPair(w http.ResponseWriter, r *http.Request, clientID, apiKey, signozURL, grantType string) {
	r = r.WithContext(util.SetSigNozURL(r.Context(), signozURL))
	accessTokenExpiresAt := time.Now().UTC().Add(h.config.AccessTokenTTL)
	accessToken, err := EncryptToken(apiKey, signozURL, clientID, accessTokenExpiresAt, h.tokenSecret)
	if err != nil {
		h.writeOAuthError(r, w, http.StatusInternalServerError, "server_error", "failed to create access token")
		return
	}

	refreshTokenValue, err := EncryptRefreshToken(
		apiKey,
		signozURL,
		clientID,
		time.Now().UTC().Add(h.config.RefreshTokenTTL),
		h.tokenSecret,
	)
	if err != nil {
		h.writeOAuthError(r, w, http.StatusInternalServerError, "server_error", "failed to create refresh token")
		return
	}

	h.writeTokenResponse(w, tokenResponse{
		AccessToken:  accessToken,
		TokenType:    "Bearer",
		ExpiresIn:    int(h.config.AccessTokenTTL.Seconds()),
		RefreshToken: refreshTokenValue,
	})

	h.emit(r.Context(), analytics.EventOAuthTokenIssued, apiKey, signozURL, map[string]any{
		analytics.AttrTenantURL: signozURL,
		analytics.AttrGrantType: grantType,
	})
}

func (h *Handler) validateAuthorizeRequest(values url.Values) (authorizeTemplateData, error) {
	clientID := strings.TrimSpace(values.Get("client_id"))
	redirectURI := strings.TrimSpace(values.Get("redirect_uri"))
	codeChallenge := strings.TrimSpace(values.Get("code_challenge"))
	codeChallengeMethod := strings.TrimSpace(values.Get("code_challenge_method"))

	if responseType := values.Get("response_type"); responseType != "" && responseType != "code" {
		return authorizeTemplateData{}, fmt.Errorf("response_type must be code")
	}
	if clientID == "" || redirectURI == "" {
		return authorizeTemplateData{}, fmt.Errorf("client_id and redirect_uri are required")
	}
	if codeChallenge == "" || codeChallengeMethod != "S256" {
		return authorizeTemplateData{}, fmt.Errorf("code_challenge and code_challenge_method=S256 are required")
	}

	redirectURIs, clientName, _, err := DecryptClientID(clientID, h.tokenSecret)
	if err != nil {
		return authorizeTemplateData{}, fmt.Errorf("client_id is not registered")
	}
	if !registeredRedirectURI(redirectURIs, redirectURI) {
		return authorizeTemplateData{}, fmt.Errorf("redirect_uri does not match the registered client")
	}

	return authorizeTemplateData{
		ClientID:            clientID,
		ClientName:          clientName,
		RedirectURI:         redirectURI,
		State:               values.Get("state"),
		CodeChallenge:       codeChallenge,
		CodeChallengeMethod: codeChallengeMethod,
		Scope:               values.Get("scope"),
	}, nil
}

func (h *Handler) validateCSRF(r *http.Request) bool {
	cookie, err := r.Cookie(csrfCookieName)
	if err != nil {
		return false
	}
	formToken := strings.TrimSpace(r.FormValue("csrf_token"))
	return formToken != "" && cookie.Value == formToken
}

func (h *Handler) writeTokenResponse(w http.ResponseWriter, resp tokenResponse) {
	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("Pragma", "no-cache")
	h.writeJSON(w, http.StatusOK, resp)
}

func (h *Handler) recordOAuthFailure(ctx context.Context, r *http.Request, status int, code, description string) {
	attrs := []slog.Attr{
		slog.Int("http.response.status_code", status),
		slog.String("oauth.error_code", code),
		slog.String("oauth.error_description", description),
	}
	attrs = append(attrs, logpkg.HTTPRequestAttrs(r)...)

	if h.meters != nil {
		metricAttrs := []attribute.KeyValue{
			attribute.String("oauth.error_code", code),
			attribute.Int("http.response.status_code", status),
		}
		metricAttrs = otelpkg.AppendTenantURL(ctx, metricAttrs)
		h.meters.OAuthFailures.Add(ctx, 1, metric.WithAttributes(metricAttrs...))
	}

	level := slog.LevelWarn
	if status >= http.StatusInternalServerError {
		level = slog.LevelError
	}
	h.logger.LogAttrs(ctx, level, "OAuth request failed", attrs...)
}

func (h *Handler) writeOAuthError(r *http.Request, w http.ResponseWriter, status int, code, description string) {
	ctx := context.Background()
	if r != nil {
		ctx = r.Context()
	}

	h.recordOAuthFailure(ctx, r, status, code, description)

	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("Pragma", "no-cache")
	h.writeJSON(w, status, map[string]string{
		"error":             code,
		"error_description": description,
	})
}

func (h *Handler) writeJSON(w http.ResponseWriter, status int, payload any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	if err := json.NewEncoder(w).Encode(payload); err != nil {
		h.logger.ErrorContext(context.Background(), "failed to write JSON response", logpkg.ErrAttr(err))
	}
}

func (h *Handler) issuerURL() string {
	return strings.TrimSuffix(h.config.OAuthIssuerURL, "/")
}

func (h *Handler) authorizePath() string {
	parsed, err := url.Parse(h.issuerURL())
	if err != nil {
		return "/oauth/authorize"
	}

	basePath := strings.TrimSuffix(parsed.Path, "/")
	if basePath == "" {
		return "/oauth/authorize"
	}

	return basePath + "/oauth/authorize"
}

// isSecure derives the Secure cookie flag from the issuer URL scheme rather
// than r.TLS, which is nil behind a TLS-terminating reverse proxy.
func (h *Handler) isSecure() bool {
	return strings.HasPrefix(h.config.OAuthIssuerURL, "https://")
}

func registeredRedirectURI(redirectURIs []string, redirectURI string) bool {
	for _, candidate := range redirectURIs {
		if candidate == redirectURI {
			return true
		}
	}
	return false
}

func validateRedirectURI(raw string) error {
	parsed, err := url.Parse(raw)
	if err != nil {
		return fmt.Errorf("redirect URI is malformed: %w", err)
	}
	if parsed.Scheme == "" {
		return fmt.Errorf("redirect URI must include a scheme")
	}
	if parsed.Fragment != "" {
		return fmt.Errorf("redirect URI fragments are not allowed")
	}

	host := strings.ToLower(parsed.Hostname())
	switch parsed.Scheme {
	case "http":
		// HTTP only allowed for loopback addresses (RFC 8252 §7.3)
		if host != "localhost" && host != "127.0.0.1" && host != "::1" {
			return fmt.Errorf("HTTP redirect URIs are only allowed for localhost, 127.0.0.1, or [::1]")
		}
	case "https":
		if host == "" {
			return fmt.Errorf("redirect URI host is required")
		}
	default:
		if !isAllowedCustomRedirectScheme(parsed.Scheme) {
			return fmt.Errorf("redirect URI scheme %q is not supported", parsed.Scheme)
		}
	}

	return nil
}

func isAllowedCustomRedirectScheme(scheme string) bool {
	switch strings.ToLower(scheme) {
	case "claude", "cursor":
		return true
	}

	// RFC 8252 private-use schemes should normally use a reverse-domain name.
	return strings.Contains(scheme, ".")
}

func addAuthorizeResponse(redirectURI, code, state string) (string, error) {
	parsed, err := url.Parse(redirectURI)
	if err != nil {
		return "", err
	}

	query := parsed.Query()
	query.Set("code", code)
	if state != "" {
		query.Set("state", state)
	}
	parsed.RawQuery = query.Encode()
	return parsed.String(), nil
}

func randomURLSafeString(size int) (string, error) {
	buf := make([]byte, size)
	if _, err := rand.Read(buf); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(buf), nil
}
