package colabauth

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/gin-contrib/sessions"
	"github.com/gin-contrib/sessions/cookie"
	"github.com/gin-gonic/gin"
	"github.com/sigbit/mcp-auth-proxy/v2/pkg/auth"
	"github.com/stretchr/testify/require"
)

const owner = "did:plc:owner"

type env struct {
	t        *testing.T
	h        *Handler
	srv      *httptest.Server
	path     string
	exchange atomic.Int32
	// provider behaviour
	refresh    string
	profile    string
	tokenCode  int
	challenge  string
	verifierOK atomic.Bool
	assign     func(http.ResponseWriter)
	probes     atomic.Int32
	probeHdr   http.Header
	probeQuery string
}

func newEnv(t *testing.T, account string) *env {
	t.Helper()
	e := &env{t: t, refresh: "rt", tokenCode: 200,
		profile: `{"sub":"1","email":"Me@Example.com","email_verified":true}`}
	e.assign = func(w http.ResponseWriter) { w.Write([]byte(`{"assignments":[]}`)) }
	provider := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/token":
			e.exchange.Add(1)
			require.NoError(t, r.ParseForm())
			sum := sha256.Sum256([]byte(r.PostForm.Get("code_verifier")))
			e.verifierOK.Store(base64.RawURLEncoding.EncodeToString(sum[:]) == e.challenge)
			if e.tokenCode != 200 {
				w.WriteHeader(e.tokenCode)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			body := map[string]any{"access_token": "at", "token_type": "Bearer", "expires_in": 3600}
			if e.refresh != "" {
				body["refresh_token"] = e.refresh
			}
			json.NewEncoder(w).Encode(body)
		case "/userinfo":
			require.Equal(t, "Bearer at", r.Header.Get("Authorization"))
			w.Write([]byte(e.profile))
		case "/assignments":
			e.probes.Add(1)
			e.probeHdr, e.probeQuery = r.Header.Clone(), r.URL.RawQuery
			e.assign(w)
		}
	}))
	t.Cleanup(provider.Close)

	e.path = filepath.Join(t.TempDir(), "creds", "colab.json")
	h, err := New("https://gw.example/", "cid", "csecret", owner, e.path, account, defaultScopes)
	require.NoError(t, err)
	h.oauth.Endpoint.TokenURL = provider.URL + "/token"
	h.userinfoURL = provider.URL + "/userinfo"
	h.probeURL = provider.URL + "/assignments?authuser=0"
	e.h = h

	gin.SetMode(gin.TestMode)
	r := gin.New()
	r.Use(sessions.Sessions("session", cookie.NewStore([]byte("0123456789abcdef0123456789abcdef"))))
	ar, err := auth.NewAuthRouter(nil, false, nil)
	require.NoError(t, err)
	r.GET("/test-login", func(c *gin.Context) {
		s := sessions.Default(c)
		s.Set(auth.SessionKeyAuthorized, true)
		s.Set(auth.SessionKeyUserID, c.Query("u"))
		s.Save()
	})
	h.SetupRoutes(r, ar)
	e.srv = httptest.NewServer(r)
	t.Cleanup(e.srv.Close)
	return e
}

func (e *env) browser(user string) *http.Client {
	jar, _ := cookiejar.New(nil)
	c := &http.Client{Jar: jar, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
	if user != "" {
		res, err := c.Get(e.srv.URL + "/test-login?u=" + url.QueryEscape(user))
		require.NoError(e.t, err)
		res.Body.Close()
	}
	return c
}

func body(t *testing.T, res *http.Response) string {
	defer res.Body.Close()
	raw, err := io.ReadAll(res.Body)
	require.NoError(t, err)
	return string(raw)
}

var csrfRe = regexp.MustCompile(`name="csrf" value="([^"]+)"`)

func (e *env) csrf(c *http.Client) string {
	res, err := c.Get(e.srv.URL + PageEndpoint)
	require.NoError(e.t, err)
	m := csrfRe.FindStringSubmatch(body(e.t, res))
	require.NotNil(e.t, m)
	return m[1]
}

func (e *env) post(c *http.Client, token, origin string) *http.Response {
	req, _ := http.NewRequest(http.MethodPost, e.srv.URL+LoginEndpoint, strings.NewReader(url.Values{"csrf": {token}}.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	if origin != "" {
		req.Header.Set("Origin", origin)
	}
	res, err := c.Do(req)
	require.NoError(e.t, err)
	return res
}

// start runs the login form and returns the google redirect.
func (e *env) start(c *http.Client) *url.URL {
	res := e.post(c, e.csrf(c), "https://gw.example")
	res.Body.Close()
	require.Equal(e.t, http.StatusFound, res.StatusCode)
	loc, err := url.Parse(res.Header.Get("Location"))
	require.NoError(e.t, err)
	e.challenge = loc.Query().Get("code_challenge")
	return loc
}

func (e *env) callback(c *http.Client, query string) *http.Response {
	res, err := c.Get(e.srv.URL + CallbackEndpoint + "?" + query)
	require.NoError(e.t, err)
	res.Body.Close()
	return res
}

func (e *env) old() {
	require.NoError(e.t, os.MkdirAll(filepath.Dir(e.path), 0o700))
	require.NoError(e.t, os.WriteFile(e.path, []byte("old"), 0o600))
}

func (e *env) kept() {
	got, err := os.ReadFile(e.path)
	require.NoError(e.t, err)
	require.Equal(e.t, "old", string(got))
}

func TestFlowWritesCompatibleCredentials(t *testing.T) {
	e := newEnv(t, "me@example.COM")
	c := e.browser(owner)
	loc := e.start(c)
	q := loc.Query()
	require.Equal(t, "accounts.google.com", loc.Host)
	require.Equal(t, "S256", q.Get("code_challenge_method"))
	require.Equal(t, "offline", q.Get("access_type"))
	require.Equal(t, "consent select_account", q.Get("prompt"))
	require.Equal(t, "https://gw.example"+CallbackEndpoint, q.Get("redirect_uri"))
	require.Equal(t, "openid https://www.googleapis.com/auth/userinfo.email", q.Get("scope"))

	res := e.callback(c, "code=abc&state="+q.Get("state")+"&scope=x")
	require.Equal(t, http.StatusFound, res.StatusCode)
	require.Equal(t, PageEndpoint, res.Header.Get("Location"))
	require.True(t, e.verifierOK.Load())

	info, err := os.Stat(e.h.accountPath())
	require.NoError(t, err)
	require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
	dir, _ := os.Stat(filepath.Dir(e.path))
	require.Equal(t, os.FileMode(0o700), dir.Mode().Perm())
	store, err := e.h.readAccounts()
	require.NoError(t, err)
	require.Len(t, store.Accounts, 1)
	raw, err := json.Marshal(store.Accounts[0])
	require.NoError(t, err)
	var got map[string]any
	require.NoError(t, json.Unmarshal(raw, &got))
	require.Regexp(t, accountIDPattern, got["account_id"])
	require.Equal(t, "1", got["google_subject"])
	require.Equal(t, "authorized_user", got["type"])
	require.Equal(t, "cid", got["client_id"])
	require.Equal(t, "csecret", got["client_secret"])
	require.Equal(t, "rt", got["refresh_token"])
	require.Equal(t, "at", got["token"])
	require.Equal(t, "Me@Example.com", got["email"])
	exp, err := time.Parse(time.RFC3339, got["expiry"].(string))
	require.NoError(t, err)
	require.WithinDuration(t, time.Now().Add(time.Hour), exp, time.Minute)

	// replay is rejected without another exchange
	res = e.callback(c, "code=abc&state="+q.Get("state"))
	require.Equal(t, http.StatusBadRequest, res.StatusCode)
	require.EqualValues(t, 1, e.exchange.Load())

	page, _ := c.Get(e.srv.URL + PageEndpoint)
	require.Contains(t, body(t, page), "connected as Me@Example.com")
}

func TestOwnerGate(t *testing.T) {
	e := newEnv(t, "")
	anon := e.browser("")
	res, _ := anon.Get(e.srv.URL + PageEndpoint)
	res.Body.Close()
	require.Equal(t, http.StatusFound, res.StatusCode)
	require.Equal(t, auth.LoginEndpoint, res.Header.Get("Location"))

	other := e.browser("did:plc:other")
	for _, path := range []string{PageEndpoint, CallbackEndpoint + "?state=x&code=y"} {
		res, _ := other.Get(e.srv.URL + path)
		res.Body.Close()
		require.Equal(t, http.StatusForbidden, res.StatusCode, path)
	}
	res = e.post(other, "x", "")
	res.Body.Close()
	require.Equal(t, http.StatusForbidden, res.StatusCode)
	require.Zero(t, e.exchange.Load())
	require.NoFileExists(t, e.h.accountPath())
}

func TestCSRFAndOrigin(t *testing.T) {
	e := newEnv(t, "")
	c := e.browser(owner)
	token := e.csrf(c)
	for name, res := range map[string]*http.Response{
		"missing token": e.post(c, "", ""),
		"wrong token":   e.post(c, "nope", ""),
		"bad origin":    e.post(c, token, "https://evil.example"),
		"null origin":   e.post(c, token, "null"),
	} {
		res.Body.Close()
		require.Equal(t, http.StatusForbidden, res.StatusCode, name)
	}
	res, _ := c.Get(e.srv.URL + LoginEndpoint)
	res.Body.Close()
	require.Equal(t, http.StatusNotFound, res.StatusCode)
	res = e.post(c, token, "")
	res.Body.Close()
	require.Equal(t, http.StatusFound, res.StatusCode)
}

func TestStateBinding(t *testing.T) {
	e := newEnv(t, "")
	c := e.browser(owner)
	state := e.start(c).Query().Get("state")

	// another owner session cannot use the state, and does not burn it
	thief := e.browser(owner)
	require.Equal(t, http.StatusBadRequest, e.callback(thief, "code=c&state="+state).StatusCode)
	require.Equal(t, http.StatusBadRequest, e.callback(c, "code=c&state=bogus").StatusCode)
	require.Equal(t, http.StatusBadRequest, e.callback(c, "code=c").StatusCode)
	require.Zero(t, e.exchange.Load())
	require.Equal(t, http.StatusFound, e.callback(c, "code=c&state="+state).StatusCode)
}

func TestExpiryAndCap(t *testing.T) {
	e := newEnv(t, "")
	now := time.Now()
	e.h.now = func() time.Time { return now }
	c := e.browser(owner)
	state := e.start(c).Query().Get("state")
	now = now.Add(pendingTTL)
	require.Equal(t, http.StatusBadRequest, e.callback(c, "code=c&state="+state).StatusCode)
	require.Zero(t, e.exchange.Load())

	for i := 0; i < maxPending; i++ {
		e.start(c)
	}
	res := e.post(c, e.csrf(c), "")
	res.Body.Close()
	require.Equal(t, http.StatusTooManyRequests, res.StatusCode)
	now = now.Add(pendingTTL)
	e.start(c)
	require.Len(t, e.h.pending, 1)
}

func TestFailuresKeepOldCredentials(t *testing.T) {
	cases := map[string]struct {
		account string
		mutate  func(*env)
		query   string
	}{
		"denied":      {query: "error=access_denied&state=%s"},
		"no code":     {query: "state=%s"},
		"token error": {mutate: func(e *env) { e.tokenCode = 400 }, query: "code=c&state=%s"},
		"no refresh":  {mutate: func(e *env) { e.refresh = "" }, query: "code=c&state=%s"},
		"unverified":  {mutate: func(e *env) { e.profile = `{"sub":"1","email":"a@b.c","email_verified":false}` }, query: "code=c&state=%s"},
		"no sub":      {mutate: func(e *env) { e.profile = `{"email":"a@b.c","email_verified":true}` }, query: "code=c&state=%s"},
		"oversized": {mutate: func(e *env) {
			e.profile = `{"sub":"1","email":"a@b.c","email_verified":true,"x":"` + strings.Repeat("a", maxProfileBytes) + `"}`
		}, query: "code=c&state=%s"},
		"wrong account": {account: "other@example.com", query: "code=c&state=%s"},
		"colab denied":  {mutate: func(e *env) { e.assign = func(w http.ResponseWriter) { w.WriteHeader(403) } }, query: "code=c&state=%s"},
		"colab error":   {mutate: func(e *env) { e.assign = func(w http.ResponseWriter) { w.WriteHeader(500) } }, query: "code=c&state=%s"},
		"colab redirect": {mutate: func(e *env) {
			e.assign = func(w http.ResponseWriter) {
				w.Header().Set("Location", "/elsewhere")
				w.WriteHeader(302)
			}
		}, query: "code=c&state=%s"},
		"colab empty":     {mutate: func(e *env) { e.assign = func(w http.ResponseWriter) {} }, query: "code=c&state=%s"},
		"colab malformed": {mutate: func(e *env) { e.assign = func(w http.ResponseWriter) { w.Write([]byte(`{"assign`)) } }, query: "code=c&state=%s"},
		"colab null":      {mutate: func(e *env) { e.assign = func(w http.ResponseWriter) { w.Write([]byte(`null`)) } }, query: "code=c&state=%s"},
		"colab wrong shape": {mutate: func(e *env) {
			e.assign = func(w http.ResponseWriter) { w.Write([]byte(`{"assignments":{"endpoint":"x"}}`)) }
		}, query: "code=c&state=%s"},
		"colab oversized": {mutate: func(e *env) {
			e.assign = func(w http.ResponseWriter) {
				w.Write([]byte(`{"assignments":[],"x":"` + strings.Repeat("a", maxProbeBytes) + `"}`))
			}
		}, query: "code=c&state=%s"},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			e := newEnv(t, tc.account)
			if tc.mutate != nil {
				tc.mutate(e)
			}
			e.old()
			c := e.browser(owner)
			state := e.start(c).Query().Get("state")
			res := e.callback(c, strings.Replace(tc.query, "%s", state, 1))
			require.GreaterOrEqual(t, res.StatusCode, 400)
			e.kept()
		})
	}
}

func TestFromEnv(t *testing.T) {
	set := func(id, secret, subject, path string) {
		t.Setenv("COLAB_GOOGLE_CLIENT_ID", id)
		t.Setenv("COLAB_GOOGLE_CLIENT_SECRET", secret)
		t.Setenv("MCP_ALLOWED_SUBJECT", subject)
		t.Setenv("COLAB_MANAGED_CREDENTIALS", path)
	}
	set("", "", "", "")
	h, err := FromEnv("https://gw.example/")
	require.NoError(t, err)
	require.Nil(t, h)

	for _, args := range [][4]string{
		{"id", "", owner, "/tmp/x"},
		{"", "secret", owner, "/tmp/x"},
		{"id", "secret", "", "/tmp/x"},
		{"id", "secret", owner, ""},
		{"id", "secret", owner, "relative.json"},
	} {
		set(args[0], args[1], args[2], args[3])
		_, err := FromEnv("https://gw.example/")
		require.Error(t, err, args)
	}

	set("id", "secret", owner, "/tmp/x.json")
	t.Setenv("COLAB_GOOGLE_SCOPES", "openid email")
	t.Setenv("COLAB_GOOGLE_ACCOUNT", " Me@Example.com ")
	h, err = FromEnv("https://gw.example/")
	require.NoError(t, err)
	require.Equal(t, []string{"openid", "email"}, h.oauth.Scopes)
	require.Equal(t, "Me@Example.com", h.account)
	require.Equal(t, "https://gw.example/.auth/colab/callback", h.oauth.RedirectURL)
}

func TestNewRejectsBadOrigins(t *testing.T) {
	for _, u := range []string{"http://gw.example", "https://u@gw.example", "https://gw.example/x", "https://gw.example/?a=1", "https://gw.example/#f", "https://gw.example/?", "gw.example", ""} {
		_, err := New(u, "id", "secret", owner, "/tmp/x", "", defaultScopes)
		require.Error(t, err, u)
	}
	for _, u := range []string{"https://gw.example", "https://gw.example/", "http://127.0.0.1:8080", "http://[::1]:8080/"} {
		_, err := New(u, "id", "secret", owner, "/tmp/x", "", defaultScopes)
		require.NoError(t, err, u)
	}
}

func TestOversizeFormLeavesNoPending(t *testing.T) {
	e := newEnv(t, "")
	c := e.browser(owner)
	token := e.csrf(c)
	res := e.post(c, token+strings.Repeat("a", maxFormBytes), "")
	res.Body.Close()
	require.Equal(t, http.StatusBadRequest, res.StatusCode)
	require.Empty(t, e.h.pending)
}

func TestPrivateHeadersEverywhere(t *testing.T) {
	e := newEnv(t, "")
	check := func(res *http.Response) {
		res.Body.Close()
		require.Equal(t, "no-store", res.Header.Get("Cache-Control"))
		require.Equal(t, "strict-origin", res.Header.Get("Referrer-Policy"))
		require.Equal(t, "nosniff", res.Header.Get("X-Content-Type-Options"))
	}
	anon := e.browser("")
	res, _ := anon.Get(e.srv.URL + PageEndpoint)
	check(res)
	c := e.browser(owner)
	res, _ = c.Get(e.srv.URL + PageEndpoint)
	check(res)
	res, _ = c.Get(e.srv.URL + CallbackEndpoint)
	check(res)
	check(e.post(c, "nope", ""))
	state := e.start(c).Query().Get("state")
	res, _ = c.Get(e.srv.URL + CallbackEndpoint + "?code=c&state=" + state)
	check(res)
}

func TestRejectsNonPrivateExistingParent(t *testing.T) {
	e := newEnv(t, "")
	e.old()
	require.NoError(t, os.Chmod(filepath.Dir(e.path), 0o755))
	c := e.browser(owner)
	state := e.start(c).Query().Get("state")
	res := e.callback(c, "code=c&state="+state)
	require.Equal(t, http.StatusInternalServerError, res.StatusCode)
	e.kept()
}

func TestColabProbeRequestAndShapes(t *testing.T) {
	for name, reply := range map[string]string{
		"object":      `{"assignments":[{"endpoint":"e"}]}`,
		"array":       `[{"endpoint":"e"}]`,
		"xssi":        ")]}'\n{\"assignments\":[]}",
		"xssi suffix": ")]}'junk\n[]",
	} {
		t.Run(name, func(t *testing.T) {
			e := newEnv(t, "")
			e.assign = func(w http.ResponseWriter) { w.Write([]byte(reply)) }
			c := e.browser(owner)
			state := e.start(c).Query().Get("state")
			require.Equal(t, http.StatusFound, e.callback(c, "code=c&state="+state).StatusCode)
			require.EqualValues(t, 1, e.probes.Load())
			require.Equal(t, "authuser=0", e.probeQuery)
			require.Equal(t, "Bearer at", e.probeHdr.Get("Authorization"))
			require.Equal(t, "application/json", e.probeHdr.Get("Accept"))
			require.Equal(t, "colab-cli", e.probeHdr.Get("X-Colab-Client-Agent"))
			require.Equal(t, "Google", e.probeHdr.Get("X-Colab-Tunnel"))
			require.FileExists(t, e.h.accountPath())
		})
	}
}

func TestNoColabProbeBeforeProfileChecks(t *testing.T) {
	for name, mutate := range map[string]func(*env){
		"wrong account": func(e *env) { e.h.account = "other@example.com" },
		"unverified":    func(e *env) { e.profile = `{"sub":"1","email":"a@b.c","email_verified":false}` },
		"no refresh":    func(e *env) { e.refresh = "" },
	} {
		t.Run(name, func(t *testing.T) {
			e := newEnv(t, "")
			mutate(e)
			c := e.browser(owner)
			state := e.start(c).Query().Get("state")
			require.GreaterOrEqual(t, e.callback(c, "code=c&state="+state).StatusCode, 400)
			require.Zero(t, e.probes.Load())
			require.NoFileExists(t, e.h.accountPath())
		})
	}
}

func TestProbeFailureHidesProviderDetail(t *testing.T) {
	e := newEnv(t, "")
	e.assign = func(w http.ResponseWriter) {
		w.WriteHeader(403)
		w.Write([]byte("secret-provider-detail"))
	}
	c := e.browser(owner)
	state := e.start(c).Query().Get("state")
	res, err := c.Get(e.srv.URL + CallbackEndpoint + "?code=c&state=" + state)
	require.NoError(t, err)
	require.Equal(t, http.StatusBadGateway, res.StatusCode)
	require.NotContains(t, body(t, res), "secret-provider-detail")
	require.NoFileExists(t, e.h.accountPath())
}

func TestProbeNetworkFailureKeepsOldCredential(t *testing.T) {
	e := newEnv(t, "")
	unavailable := httptest.NewServer(http.NotFoundHandler())
	unavailable.Close()
	e.h.probeURL = unavailable.URL
	e.old()
	c := e.browser(owner)
	state := e.start(c).Query().Get("state")
	require.Equal(t, http.StatusBadGateway, e.callback(c, "code=c&state="+state).StatusCode)
	e.kept()
}

func TestProbeDoesNotFollowRedirects(t *testing.T) {
	var followed atomic.Int32
	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		followed.Add(1)
		w.Write([]byte(`{"assignments":[]}`))
	}))
	defer target.Close()
	e := newEnv(t, "")
	e.assign = func(w http.ResponseWriter) {
		w.Header().Set("Location", target.URL)
		w.WriteHeader(http.StatusFound)
	}
	e.old()
	c := e.browser(owner)
	state := e.start(c).Query().Get("state")
	require.Equal(t, http.StatusBadGateway, e.callback(c, "code=c&state="+state).StatusCode)
	require.Zero(t, followed.Load())
	e.kept()
}