package proxy

import (
	"crypto/rsa"
	"fmt"
	"net/http"
	"os"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/golang-jwt/jwt/v5"
	"github.com/mattn/go-jsonpointer"
	"github.com/sigbit/mcp-auth-proxy/v2/pkg/utils"
)

type ProxyRouter struct {
	externalURL                string
	allowedSubject             string
	proxy                      http.Handler
	publicKey                  *rsa.PublicKey
	proxyHeaders               http.Header
	httpStreamingOnly          bool
	forwardAuthorizationHeader bool
	headerMapping              map[string]string
	headerMappingBase          string
}

func NewProxyRouter(
	externalURL string,
	proxy http.Handler,
	publicKey *rsa.PublicKey,
	proxyHeaders http.Header,
	httpStreamingOnly bool,
	forwardAuthorizationHeader bool,
	headerMapping map[string]string,
	headerMappingBase string,
) (*ProxyRouter, error) {
	return &ProxyRouter{
		externalURL:                externalURL,
		allowedSubject:             os.Getenv("MCP_ALLOWED_SUBJECT"),
		proxy:                      proxy,
		publicKey:                  publicKey,
		proxyHeaders:               proxyHeaders,
		httpStreamingOnly:          httpStreamingOnly,
		forwardAuthorizationHeader: forwardAuthorizationHeader,
		headerMapping:              headerMapping,
		headerMappingBase:          headerMappingBase,
	}, nil
}

const (
	OauthProtectedResourceEndpoint = "/.well-known/oauth-protected-resource"
)

func (p *ProxyRouter) SetupRoutes(router gin.IRouter) {
	router.GET(OauthProtectedResourceEndpoint, p.handleProtectedResource)
	router.Use(p.handleProxy)
}

type protectedResourceResponse struct {
	Resource             string   `json:"resource"`
	AuthorizationServers []string `json:"authorization_servers"`
}

func (p *ProxyRouter) handleProtectedResource(c *gin.Context) {
	c.JSON(http.StatusOK, protectedResourceResponse{
		Resource:             utils.MCPResource(p.externalURL),
		AuthorizationServers: []string{p.externalURL},
	})
}

func (p *ProxyRouter) handleProxy(c *gin.Context) {
	challenge := fmt.Sprintf(`Bearer resource_metadata="%s%s"`, strings.TrimRight(p.externalURL, "/"), OauthProtectedResourceEndpoint)
	authHeader := c.Request.Header.Get("Authorization")
	if !strings.HasPrefix(authHeader, "Bearer ") {
		c.Header("WWW-Authenticate", challenge)
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
		return
	}
	bearerToken := strings.TrimPrefix(authHeader, "Bearer ")

	claims := jwt.MapClaims{}
	token, err := jwt.ParseWithClaims(bearerToken, claims, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return p.publicKey, nil
	}, jwt.WithIssuer(p.externalURL), jwt.WithAudience(utils.MCPResource(p.externalURL)), jwt.WithExpirationRequired(), jwt.WithValidMethods([]string{"RS256"}))

	subject, _ := claims.GetSubject()
	if err != nil || !token.Valid || (p.allowedSubject != "" && subject != p.allowedSubject) {
		c.Header("WWW-Authenticate", challenge)
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "Invalid token"})
		return
	}

	if p.httpStreamingOnly && isSSEGetRequest(c.Request) {
		c.AbortWithStatusJSON(http.StatusMethodNotAllowed, gin.H{"error": "SSE (GET) streaming is not supported by this backend; use POST-based HTTP streaming instead"})
		return
	}

	if !p.forwardAuthorizationHeader {
		c.Request.Header.Del("Authorization")
	}
	for _, headerName := range p.headerMapping {
		c.Request.Header.Del(headerName)
	}
	for key, values := range p.proxyHeaders {
		if strings.EqualFold(key, "Authorization") {
			c.Request.Header.Del("Authorization")
		}
		for _, value := range values {
			c.Request.Header.Add(key, value)
		}
	}

	if len(p.headerMapping) > 0 {
		var source any = map[string]any(claims)
		if p.headerMappingBase != "/" {
			val, err := jsonpointer.Get(source, p.headerMappingBase)
			if err != nil {
				source = nil
			} else {
				source = val
			}
		}
		if source != nil {
			for pointer, headerName := range p.headerMapping {
				val, err := jsonpointer.Get(source, pointer)
				if err != nil {
					continue
				}
				switch v := val.(type) {
				case string:
					c.Request.Header.Set(headerName, v)
				case []any:
					var parts []string
					for _, item := range v {
						if s, ok := item.(string); ok {
							parts = append(parts, s)
						}
					}
					c.Request.Header.Set(headerName, strings.Join(parts, ","))
				default:
					c.Request.Header.Set(headerName, fmt.Sprintf("%v", v))
				}
			}
		}
	}

	p.proxy.ServeHTTP(c.Writer, c.Request)
}

func isSSEGetRequest(r *http.Request) bool {
	if r.Method != http.MethodGet {
		return false
	}
	accept := r.Header.Get("Accept")
	if accept == "" {
		return false
	}
	for _, value := range strings.Split(accept, ",") {
		mediaType := strings.TrimSpace(strings.ToLower(value))
		if idx := strings.Index(mediaType, ";"); idx != -1 {
			mediaType = strings.TrimSpace(mediaType[:idx])
		}
		if mediaType == "text/event-stream" {
			return true
		}
	}
	return false
}