package colabauth

import (
	"context"
	"crypto/subtle"
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/gin-contrib/sessions"
	"github.com/gin-gonic/gin"
	"github.com/sigbit/mcp-auth-proxy/v2/pkg/auth"
	"github.com/sigbit/mcp-auth-proxy/v2/pkg/utils"
	"golang.org/x/oauth2"
)

const (
	PageEndpoint     = "/.auth/colab"
	LoginEndpoint    = "/.auth/colab/login"
	CallbackEndpoint = "/.auth/colab/callback"

	authURL     = "https://accounts.google.com/o/oauth2/v2/auth"
	tokenURL    = "https://oauth2.googleapis.com/token"
	userinfoURL = "https://openidconnect.googleapis.com/v1/userinfo"
	probeURL    = "https://colab.research.google.com/tun/m/assignments?authuser=0"

	pendingTTL      = 10 * time.Minute
	maxPending      = 32
	maxProfileBytes = 8 << 10
	maxProbeBytes   = 1 << 20
	maxFormBytes    = 4 << 10
	requestTimeout  = 15 * time.Second

	sessionCSRF  = "colab_csrf"
	sessionNonce = "colab_nonce"
)

var xssiPrefix = regexp.MustCompile(`^\)\]\}'[^\n]*\n`)

var defaultScopes = []string{"openid", "https://www.googleapis.com/auth/userinfo.email"}

type pending struct {
	nonce    string
	verifier string
	expires  time.Time
}

type Handler struct {
	oauth       oauth2.Config
	subject     string
	account     string
	managedPath string
	origin      string
	userinfoURL string
	probeURL    string
	client      *http.Client
	now         func() time.Time

	storeMu sync.Mutex
	mu      sync.Mutex
	pending map[string]pending
}

func validOrigin(u *url.URL) bool {
	if u.Host == "" || u.User != nil || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || (u.Path != "" && u.Path != "/") {
		return false
	}
	return u.Scheme == "https" || (u.Scheme == "http" && (u.Hostname() == "127.0.0.1" || u.Hostname() == "::1"))
}

// FromEnv returns nil when linking is not configured.
func FromEnv(externalURL string) (*Handler, error) {
	id, secret := os.Getenv("COLAB_GOOGLE_CLIENT_ID"), os.Getenv("COLAB_GOOGLE_CLIENT_SECRET")
	if id == "" && secret == "" {
		return nil, nil
	}
	if id == "" || secret == "" {
		return nil, errors.New("COLAB_GOOGLE_CLIENT_ID and COLAB_GOOGLE_CLIENT_SECRET must be set together")
	}
	subject, managed := os.Getenv("MCP_ALLOWED_SUBJECT"), os.Getenv("COLAB_MANAGED_CREDENTIALS")
	if subject == "" {
		return nil, errors.New("MCP_ALLOWED_SUBJECT is required for colab account linking")
	}
	if !filepath.IsAbs(managed) {
		return nil, errors.New("COLAB_MANAGED_CREDENTIALS must be an absolute path for colab account linking")
	}
	scopes := strings.Fields(os.Getenv("COLAB_GOOGLE_SCOPES"))
	if len(scopes) == 0 {
		scopes = defaultScopes
	}
	h, err := New(externalURL, id, secret, subject, managed, strings.TrimSpace(os.Getenv("COLAB_GOOGLE_ACCOUNT")), scopes)
	if err != nil {
		return nil, err
	}
	if err := h.initializeAccounts(); err != nil {
		return nil, errors.New("could not initialize colab account store")
	}
	return h, nil
}

func New(externalURL, clientID, clientSecret, subject, managedPath, account string, scopes []string) (*Handler, error) {
	u, err := url.Parse(externalURL)
	if err != nil || !validOrigin(u) {
		return nil, errors.New("invalid external URL for colab account linking")
	}
	origin := u.Scheme + "://" + u.Host
	return &Handler{
		oauth: oauth2.Config{
			ClientID:     clientID,
			ClientSecret: clientSecret,
			Endpoint:     oauth2.Endpoint{AuthURL: authURL, TokenURL: tokenURL, AuthStyle: oauth2.AuthStyleInParams},
			RedirectURL:  origin + CallbackEndpoint,
			Scopes:       scopes,
		},
		subject:     subject,
		account:     account,
		managedPath: managedPath,
		origin:      origin,
		userinfoURL: userinfoURL,
		probeURL:    probeURL,
		client: &http.Client{
			Timeout:       requestTimeout,
			CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
		},
		now:     time.Now,
		pending: map[string]pending{},
	}, nil
}

func (h *Handler) SetupRoutes(router gin.IRouter, authRouter *auth.AuthRouter) {
	g := router.Group(PageEndpoint, privateHeaders, authRouter.RequireAuth(), h.requireOwner)
	g.GET("", h.handlePage)
	g.POST("/login", h.handleStart)
	g.POST("/remove", h.handleRemove)
	g.GET("/callback", h.handleCallback)
}

func privateHeaders(c *gin.Context) {
	c.Header("Cache-Control", "no-store")
	// no-referrer turns the Origin of native form posts into null.
	c.Header("Referrer-Policy", "strict-origin")
	c.Header("X-Content-Type-Options", "nosniff")
}

func (h *Handler) requireOwner(c *gin.Context) {
	user, _ := sessions.Default(c).Get(auth.SessionKeyUserID).(string)
	if user == "" || subtle.ConstantTimeCompare([]byte(user), []byte(h.subject)) != 1 {
		h.fail(c, http.StatusForbidden, "forbidden")
		return
	}
	c.Next()
}

var pageTemplate = template.Must(template.New("page").Parse(`<!doctype html>
<meta charset="utf-8"><title>colab accounts</title>
<h1>colab accounts</h1>
{{if .Accounts}}<ul>{{range .Accounts}}<li>connected as {{.Email}} <code>{{.AccountID}}</code>
<form method="post" action="/.auth/colab/remove">
<input type="hidden" name="csrf" value="{{$.CSRF}}">
<input type="hidden" name="account_id" value="{{.AccountID}}">
<button>remove from pool</button></form></li>{{end}}</ul>
<p>removing an account does not stop its runtimes or revoke google access.</p>
{{else}}<p>not connected.</p>{{end}}
<form method="post" action="/.auth/colab/login">
<input type="hidden" name="csrf" value="{{.CSRF}}">
<button>connect google account</button>
</form>
<p>connect another account to add it; reconnecting the same account keeps its id. up to 16 accounts.</p>
`))

func (h *Handler) handlePage(c *gin.Context) {
	session := sessions.Default(c)
	token, err := utils.GenerateState()
	if err != nil {
		h.fail(c, http.StatusInternalServerError, "something went wrong")
		return
	}
	session.Set(sessionCSRF, token)
	if err := session.Save(); err != nil {
		h.fail(c, http.StatusInternalServerError, "something went wrong")
		return
	}
	store, err := h.readAccounts()
	if err != nil {
		h.fail(c, http.StatusInternalServerError, "could not read accounts")
		return
	}
	c.Header("Content-Type", "text/html; charset=utf-8")
	pageTemplate.Execute(c.Writer, map[string]any{"Accounts": store.Accounts, "CSRF": token})
}

func (h *Handler) validForm(c *gin.Context) bool {
	session := sessions.Default(c)
	want, _ := session.Get(sessionCSRF).(string)
	if origin := c.GetHeader("Origin"); origin != "" && origin != h.origin {
		h.fail(c, http.StatusForbidden, "forbidden")
		return false
	}
	c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, maxFormBytes)
	if err := c.Request.ParseForm(); err != nil {
		h.fail(c, http.StatusBadRequest, "invalid form")
		return false
	}
	got := c.Request.PostForm.Get("csrf")
	if want == "" || subtle.ConstantTimeCompare([]byte(got), []byte(want)) != 1 {
		h.fail(c, http.StatusForbidden, "forbidden")
		return false
	}
	return true
}

func (h *Handler) handleRemove(c *gin.Context) {
	if !h.validForm(c) {
		return
	}
	removed, err := h.removeAccount(c.Request.PostForm.Get("account_id"))
	if err != nil {
		h.fail(c, http.StatusInternalServerError, "could not remove account")
		return
	}
	if !removed {
		h.fail(c, http.StatusNotFound, "account not found")
		return
	}
	c.Redirect(http.StatusSeeOther, PageEndpoint)
}

func (h *Handler) handleStart(c *gin.Context) {
	if !h.validForm(c) {
		return
	}
	session := sessions.Default(c)
	state, err := utils.GenerateState()
	nonce, err2 := utils.GenerateState()
	if err != nil || err2 != nil {
		h.fail(c, http.StatusInternalServerError, "something went wrong")
		return
	}
	verifier := oauth2.GenerateVerifier()
	now := h.now()
	h.mu.Lock()
	h.prune(now)
	if len(h.pending) >= maxPending {
		h.mu.Unlock()
		h.fail(c, http.StatusTooManyRequests, "too many pending links, try again later")
		return
	}
	h.pending[state] = pending{nonce: nonce, verifier: verifier, expires: now.Add(pendingTTL)}
	h.mu.Unlock()

	session.Set(sessionNonce, nonce)
	if err := session.Save(); err != nil {
		h.take(state, nonce)
		h.fail(c, http.StatusInternalServerError, "something went wrong")
		return
	}
	c.Redirect(http.StatusFound, h.oauth.AuthCodeURL(state,
		oauth2.AccessTypeOffline,
		oauth2.SetAuthURLParam("prompt", "consent select_account"),
		oauth2.S256ChallengeOption(verifier)))
}

func (h *Handler) prune(now time.Time) {
	for k, p := range h.pending {
		if !now.Before(p.expires) {
			delete(h.pending, k)
		}
	}
}

// take consumes the pending entry only when it is live and bound to nonce.
func (h *Handler) take(state, nonce string) (pending, bool) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.prune(h.now())
	p, ok := h.pending[state]
	if !ok || nonce == "" || subtle.ConstantTimeCompare([]byte(p.nonce), []byte(nonce)) != 1 {
		return pending{}, false
	}
	delete(h.pending, state)
	return p, true
}

func (h *Handler) handleCallback(c *gin.Context) {
	session := sessions.Default(c)
	nonce, _ := session.Get(sessionNonce).(string)
	state := c.Query("state")
	if state == "" {
		h.fail(c, http.StatusBadRequest, "invalid link request")
		return
	}
	p, ok := h.take(state, nonce)
	if !ok {
		h.fail(c, http.StatusBadRequest, "invalid link request")
		return
	}
	session.Delete(sessionNonce)
	if err := session.Save(); err != nil {
		h.fail(c, http.StatusInternalServerError, "something went wrong")
		return
	}
	code := c.Query("code")
	if c.Query("error") != "" || code == "" {
		h.fail(c, http.StatusBadRequest, "google did not grant access")
		return
	}
	status, err := h.link(c.Request.Context(), code, p.verifier)
	if err != nil {
		h.fail(c, status, err.Error())
		return
	}
	c.Redirect(http.StatusFound, PageEndpoint)
}

type credentials struct {
	AccountID     string   `json:"account_id,omitempty"`
	GoogleSubject string   `json:"google_subject,omitempty"`
	Type          string   `json:"type"`
	ClientID      string   `json:"client_id"`
	ClientSecret  string   `json:"client_secret"`
	RefreshToken  string   `json:"refresh_token"`
	Token         string   `json:"token"`
	Expiry        string   `json:"expiry"`
	Email         string   `json:"email,omitempty"`
	Scopes        []string `json:"scopes,omitempty"`
}

// link exchanges the code and persists credentials only after every check passes.
func (h *Handler) link(ctx context.Context, code, verifier string) (int, error) {
	ctx, cancel := context.WithTimeout(ctx, requestTimeout)
	defer cancel()
	ctx = context.WithValue(ctx, oauth2.HTTPClient, h.client)
	token, err := h.oauth.Exchange(ctx, code, oauth2.VerifierOption(verifier))
	if err != nil {
		return http.StatusBadGateway, errors.New("google token exchange failed")
	}
	if token.RefreshToken == "" {
		return http.StatusBadGateway, errors.New("google returned no refresh token")
	}
	email, googleSubject, err := h.profileIdentity(ctx, token.AccessToken)
	if err != nil {
		return http.StatusBadGateway, errors.New("could not verify google account")
	}
	if h.account != "" && !strings.EqualFold(email, h.account) {
		return http.StatusForbidden, errors.New("google account is not the configured account")
	}
	if err := h.probeColab(ctx, token.AccessToken); err != nil {
		return http.StatusBadGateway, errors.New("could not verify colab access for this google client or account")
	}
	creds := credentials{
		Type:          "authorized_user",
		ClientID:      h.oauth.ClientID,
		ClientSecret:  h.oauth.ClientSecret,
		RefreshToken:  token.RefreshToken,
		Token:         token.AccessToken,
		Expiry:        token.Expiry.UTC().Format(time.RFC3339),
		Email:         email,
		GoogleSubject: googleSubject,
		Scopes:        h.oauth.Scopes,
	}
	if err := h.save(creds); err != nil {
		if errors.Is(err, errPoolFull) {
			return http.StatusConflict, errPoolFull
		}
		return http.StatusInternalServerError, errors.New("could not save credentials")
	}
	return http.StatusOK, nil
}

func (h *Handler) profileIdentity(ctx context.Context, accessToken string) (string, string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, h.userinfoURL, nil)
	if err != nil {
		return "", "", err
	}
	req.Header.Set("Authorization", "Bearer "+accessToken)
	res, err := h.client.Do(req)
	if err != nil {
		return "", "", err
	}
	defer res.Body.Close()
	body, err := io.ReadAll(io.LimitReader(res.Body, maxProfileBytes+1))
	if err != nil || res.StatusCode != http.StatusOK || len(body) > maxProfileBytes {
		return "", "", errors.New("bad profile response")
	}
	var profile struct {
		Sub           string `json:"sub"`
		Email         string `json:"email"`
		EmailVerified bool   `json:"email_verified"`
	}
	if json.Unmarshal(body, &profile) != nil || profile.Sub == "" || profile.Email == "" || !profile.EmailVerified {
		return "", "", errors.New("unverified profile")
	}
	return profile.Email, profile.Sub, nil
}

func (h *Handler) probeColab(ctx context.Context, accessToken string) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, h.probeURL, nil)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+accessToken)
	req.Header.Set("Accept", "application/json")
	req.Header.Set("X-Colab-Client-Agent", "colab-cli")
	req.Header.Set("X-Colab-Tunnel", "Google")
	res, err := h.client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	body, err := io.ReadAll(io.LimitReader(res.Body, maxProbeBytes+1))
	if err != nil || res.StatusCode != http.StatusOK || len(body) > maxProbeBytes {
		return errors.New("bad colab response")
	}
	body = xssiPrefix.ReplaceAll(body, nil)
	var data any
	if err := json.Unmarshal(body, &data); err != nil {
		return err
	}
	if obj, ok := data.(map[string]any); ok {
		data = obj["assignments"]
	}
	if _, ok := data.([]any); !ok {
		return errors.New("invalid assignments list")
	}
	return nil
}

func ensurePrivateDir(dir string) error {
	info, err := os.Stat(dir)
	if errors.Is(err, os.ErrNotExist) {
		if err := os.MkdirAll(dir, 0o700); err != nil {
			return err
		}
		return os.Chmod(dir, 0o700)
	}
	if err != nil {
		return err
	}
	if !info.IsDir() || info.Mode().Perm()&0o077 != 0 {
		return errors.New("credential directory is not private")
	}
	return nil
}

func (h *Handler) fail(c *gin.Context, status int, message string) {
	page := fmt.Sprintf("<!doctype html><title>colab account</title><p>%s</p>", template.HTMLEscapeString(message))
	c.Data(status, "text/html; charset=utf-8", []byte(page))
	c.Abort()
}