package identityprovider

import (
	"bytes"
	"context"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"errors"
	"io"
	"math/big"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"vpn-platform/backend/internal/model"
	"vpn-platform/backend/internal/store"
)

type requestState struct {
	Redirect  string `json:"redirect"`
	ClientID  string `json:"client_id"`
	State     string `json:"state"`
	Nonce     string `json:"nonce"`
	Challenge string `json:"challenge"`
}
type codeData struct {
	requestState
	Subject string `json:"subject"`
	Name    string `json:"name"`
}
type accessData struct {
	Subject string `json:"sub"`
	Name    string `json:"name"`
}

func noStore(c *gin.Context) { c.Header("Cache-Control", "no-store"); c.Header("Pragma", "no-cache") }
func oidcError(c *gin.Context, status int) {
	noStore(c)
	c.JSON(status, gin.H{"error": "invalid_request"})
}

func (p *Provider) Discovery(c *gin.Context) {
	_, base, err := enabled()
	if err != nil {
		oidcError(c, http.StatusNotFound)
		return
	}
	noStore(c)
	c.JSON(200, gin.H{
		"issuer": base, "authorization_endpoint": base + "/authorize", "token_endpoint": base + "/token",
		"userinfo_endpoint": base + "/userinfo", "jwks_uri": base + "/jwks",
		"response_types_supported": []string{"code"}, "subject_types_supported": []string{"public"},
		"id_token_signing_alg_values_supported": []string{"RS256"}, "scopes_supported": []string{"openid"},
		"grant_types_supported":                 []string{"authorization_code"},
		"token_endpoint_auth_methods_supported": []string{"client_secret_basic", "client_secret_post"},
		"code_challenge_methods_supported":      []string{"S256"},
	})
}

func (p *Provider) JWKS(c *gin.Context) {
	if _, _, err := enabled(); err != nil {
		oidcError(c, 404)
		return
	}
	key, kid, err := signingKey()
	if err != nil {
		oidcError(c, 503)
		return
	}
	noStore(c)
	c.JSON(200, gin.H{"keys": []gin.H{{"kty": "RSA", "use": "sig", "alg": "RS256", "kid": kid,
		"n": base64.RawURLEncoding.EncodeToString(key.N.Bytes()), "e": base64.RawURLEncoding.EncodeToString(big.NewInt(int64(key.E)).Bytes())}}})
}

func (p *Provider) Authorize(c *gin.Context) {
	cfg, base, err := enabled()
	if err != nil {
		oidcError(c, 503)
		return
	}
	redirect := c.Query("redirect_uri")
	domain := strings.TrimSuffix(strings.TrimPrefix(base, "https://"), "/id")
	if !validateRedirect(redirect, domain) || c.Query("client_id") != cfg.ClientID || c.Query("response_type") != "code" || !strings.Contains(" "+c.Query("scope")+" ", " openid ") {
		oidcError(c, 400)
		return
	}
	challenge, method := c.Query("code_challenge"), c.Query("code_challenge_method")
	if (challenge != "" && (method != "S256" || len(challenge) != 43)) || (challenge == "" && method != "") || len(c.Query("state")) > 1024 || len(c.Query("nonce")) > 1024 {
		oidcError(c, 400)
		return
	}
	state, err := randomToken()
	if err != nil {
		oidcError(c, 503)
		return
	}
	pending := requestState{Redirect: redirect, ClientID: cfg.ClientID, State: c.Query("state"), Nonce: c.Query("nonce"), Challenge: challenge}
	if err := saveGrant("state", state, pending, 5*time.Minute); err != nil {
		oidcError(c, 503)
		return
	}
	upstream := feishuAuthorize + "?" + url.Values{
		"client_id": {cfg.AppID}, "redirect_uri": {base + "/feishu/callback"}, "response_type": {"code"}, "scope": {"user_profile"}, "state": {state},
	}.Encode()
	noStore(c)
	c.Redirect(http.StatusFound, upstream)
}

var upstreamClient = &http.Client{Timeout: 8 * time.Second, CheckRedirect: func(_ *http.Request, _ []*http.Request) error { return http.ErrUseLastResponse }}

func postFeishuToken(ctx context.Context, appID, appSecret, code, callback string) (string, error) {
	payload, _ := json.Marshal(gin.H{"grant_type": "authorization_code", "client_id": appID, "client_secret": appSecret, "code": code, "redirect_uri": callback})
	req, err := http.NewRequestWithContext(ctx, "POST", feishuToken, bytes.NewReader(payload))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	response, err := upstreamClient.Do(req)
	if err != nil {
		return "", err
	}
	defer response.Body.Close()
	if response.StatusCode != 200 {
		return "", errors.New("upstream token rejected")
	}
	var data struct {
		Code        *int   `json:"code"`
		AccessToken string `json:"access_token"`
		Data        struct {
			AccessToken string `json:"access_token"`
		} `json:"data"`
	}
	if err := json.NewDecoder(io.LimitReader(response.Body, 1<<20)).Decode(&data); err != nil || data.Code == nil || *data.Code != 0 {
		return "", errors.New("upstream token invalid")
	}
	if data.AccessToken != "" {
		return data.AccessToken, nil
	}
	if data.Data.AccessToken != "" {
		return data.Data.AccessToken, nil
	}
	return "", errors.New("upstream access token missing")
}

func fetchFeishuUser(ctx context.Context, token string) (accessData, string, error) {
	req, err := http.NewRequestWithContext(ctx, "GET", feishuUserInfo, nil)
	if err != nil {
		return accessData{}, "", err
	}
	req.Header.Set("Authorization", "Bearer "+token)
	response, err := upstreamClient.Do(req)
	if err != nil {
		return accessData{}, "", err
	}
	defer response.Body.Close()
	if response.StatusCode != 200 {
		return accessData{}, "", errors.New("upstream user rejected")
	}
	var result struct {
		Code *int `json:"code"`
		Data struct {
			OpenID    string `json:"open_id"`
			TenantKey string `json:"tenant_key"`
			Name      string `json:"name"`
		} `json:"data"`
	}
	if err := json.NewDecoder(io.LimitReader(response.Body, 1<<20)).Decode(&result); err != nil || result.Code == nil || *result.Code != 0 || result.Data.OpenID == "" || result.Data.TenantKey == "" {
		return accessData{}, "", errors.New("upstream user missing stable identity")
	}
	return accessData{Subject: result.Data.TenantKey + ":" + result.Data.OpenID, Name: result.Data.Name}, result.Data.TenantKey, nil
}

func (p *Provider) FeishuCallback(c *gin.Context) {
	cfg, base, err := enabled()
	if err != nil {
		oidcError(c, 503)
		return
	}
	var state requestState
	if err := consumeGrant("state", c.Query("state"), &state); err != nil || c.Query("error") != "" || c.Query("code") == "" || len(c.Query("code")) > 2048 {
		oidcError(c, 400)
		return
	}
	secret, err := store.DecryptSecret(cfg.AppSecret)
	if err != nil {
		oidcError(c, 503)
		return
	}
	token, err := postFeishuToken(c.Request.Context(), cfg.AppID, secret, c.Query("code"), base+"/feishu/callback")
	if err != nil {
		oidcError(c, 502)
		return
	}
	user, tenant, err := fetchFeishuUser(c.Request.Context(), token)
	if err != nil {
		oidcError(c, 502)
		return
	}
	if tenant != cfg.TenantKey {
		oidcError(c, 403)
		return
	}
	code, err := randomToken()
	if err != nil {
		oidcError(c, 503)
		return
	}
	grant := codeData{requestState: state, Subject: user.Subject, Name: user.Name}
	if err := saveGrant("code", code, grant, 2*time.Minute); err != nil {
		oidcError(c, 503)
		return
	}
	u, _ := url.Parse(state.Redirect)
	q := u.Query()
	q.Set("code", code)
	if state.State != "" {
		q.Set("state", state.State)
	}
	u.RawQuery = q.Encode()
	noStore(c)
	c.Redirect(http.StatusFound, u.String())
}

func bearer(c *gin.Context, cfg model.IdentityProvider) bool {
	id, secret, ok := c.Request.BasicAuth()
	if !ok {
		id, secret = c.PostForm("client_id"), c.PostForm("client_secret")
	}
	stored, err := store.DecryptSecret(cfg.ClientSecret)
	return err == nil && subtle.ConstantTimeCompare([]byte(id), []byte(cfg.ClientID)) == 1 && subtle.ConstantTimeCompare([]byte(secret), []byte(stored)) == 1
}

func signIDToken(key *rsa.PrivateKey, kid, base string, cfg model.IdentityProvider, grant codeData) (string, error) {
	now := time.Now().Unix()
	header, _ := json.Marshal(gin.H{"alg": "RS256", "typ": "JWT", "kid": kid})
	claims := gin.H{"iss": base, "aud": cfg.ClientID, "sub": grant.Subject, "iat": now, "exp": now + 600, "auth_time": now}
	if grant.Nonce != "" {
		claims["nonce"] = grant.Nonce
	}
	body, _ := json.Marshal(claims)
	unsigned := base64.RawURLEncoding.EncodeToString(header) + "." + base64.RawURLEncoding.EncodeToString(body)
	hash := sha256.Sum256([]byte(unsigned))
	signature, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, hash[:])
	if err != nil {
		return "", err
	}
	return unsigned + "." + base64.RawURLEncoding.EncodeToString(signature), nil
}

func (p *Provider) Token(c *gin.Context) {
	cfg, base, err := enabled()
	if err != nil {
		oidcError(c, 503)
		return
	}
	noStore(c)
	if !bearer(c, cfg) || c.PostForm("grant_type") != "authorization_code" {
		c.JSON(401, gin.H{"error": "invalid_client"})
		return
	}
	var grant codeData
	if err := consumeGrant("code", c.PostForm("code"), &grant); err != nil {
		c.JSON(400, gin.H{"error": "invalid_grant"})
		return
	}
	if grant.ClientID != cfg.ClientID || grant.Redirect != c.PostForm("redirect_uri") {
		c.JSON(400, gin.H{"error": "invalid_grant"})
		return
	}
	if grant.Challenge != "" {
		verifier := c.PostForm("code_verifier")
		if len(verifier) < 43 || len(verifier) > 128 {
			c.JSON(400, gin.H{"error": "invalid_grant"})
			return
		}
		hash := sha256.Sum256([]byte(verifier))
		if subtle.ConstantTimeCompare([]byte(grant.Challenge), []byte(base64.RawURLEncoding.EncodeToString(hash[:]))) != 1 {
			c.JSON(400, gin.H{"error": "invalid_grant"})
			return
		}
	}
	key, kid, err := signingKey()
	if err != nil {
		oidcError(c, 503)
		return
	}
	idToken, err := signIDToken(key, kid, base, cfg, grant)
	if err != nil {
		oidcError(c, 503)
		return
	}
	access, err := randomToken()
	if err != nil || saveGrant("access", access, accessData{Subject: grant.Subject, Name: grant.Name}, 10*time.Minute) != nil {
		oidcError(c, 503)
		return
	}
	c.JSON(200, gin.H{"access_token": access, "id_token": idToken, "token_type": "Bearer", "expires_in": 600, "scope": "openid"})
}

func (p *Provider) UserInfo(c *gin.Context) {
	if _, _, err := enabled(); err != nil {
		oidcError(c, 503)
		return
	}
	token := strings.TrimPrefix(c.GetHeader("Authorization"), "Bearer ")
	if token == c.GetHeader("Authorization") || len(token) < 32 {
		oidcError(c, 401)
		return
	}
	var grant model.IdentityGrant
	if err := store.DB.Where("kind = ? AND digest = ? AND expires_at > ?", "access", digest(token), time.Now()).First(&grant).Error; err != nil {
		oidcError(c, 401)
		return
	}
	var user accessData
	if json.Unmarshal([]byte(grant.Payload), &user) != nil {
		oidcError(c, 401)
		return
	}
	noStore(c)
	c.JSON(200, gin.H{"sub": user.Subject, "name": user.Name})
}