// Package nativeauth provides proof-bound browser approval for native clients.
// Apps retain their existing session, admission and resource authorization rules.
package nativeauth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"html/template"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/carlosframework/rastrillo/csrf"
	"github.com/carlosframework/rastrillo/sessions"
)

type Auth struct {
	mu           sync.Mutex
	pending      map[string]*request
	sessions     *sessions.Sessions
	require      func(http.Handler) http.Handler
	origin, name string
}
type request struct {
	challenge, cookie, callback string
	expires                     time.Time
}

func New(s *sessions.Sessions, origin, name string, require func(http.Handler) http.Handler) *Auth {
	return &Auth{pending: map[string]*request{}, sessions: s, origin: origin, name: name, require: require}
}
func token() string {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	return base64.RawURLEncoding.EncodeToString(b)
}
func valid(s string) bool {
	b, e := base64.RawURLEncoding.DecodeString(s)
	return e == nil && len(b) == 32 && base64.RawURLEncoding.EncodeToString(b) == s
}
func code(id string) string {
	b := sha256.Sum256([]byte(id))
	s := strings.ToUpper(hex.EncodeToString(b[:4]))
	return s[:4] + "-" + s[4:]
}
func (a *Auth) find(id string) *request {
	for key, p := range a.pending {
		if !time.Now().Before(p.expires) {
			delete(a.pending, key)
		}
	}
	return a.pending[id]
}
func (a *Auth) pendingName() string {
	if strings.HasPrefix(a.origin, "https:") {
		return "__Host-native_pending"
	}
	return "native_pending"
}
func (a *Auth) Resume(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == "GET" && r.URL.Path == "/" {
			if c, e := r.Cookie(a.pendingName()); e == nil && valid(c.Value) {
				if _, ok := a.sessions.From(r); ok {
					http.SetCookie(w, &http.Cookie{Name: a.pendingName(), Path: "/", MaxAge: -1, Secure: strings.HasPrefix(a.origin, "https:"), HttpOnly: true, SameSite: http.SameSiteLaxMode})
					http.Redirect(w, r, "/native/authorize?id="+c.Value, 303)
					return
				}
			}
		}
		next.ServeHTTP(w, r)
	})
}
func (a *Auth) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("Referrer-Policy", "no-referrer")
	csrf.Protect(a.origin)(http.HandlerFunc(a.route)).ServeHTTP(w, r)
}
func (a *Auth) route(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "POST" && r.URL.Path == "/native/request":
		a.begin(w, r)
	case r.Method == "POST" && r.URL.Path == "/native/exchange":
		a.exchange(w, r)
	case r.Method == "GET" && r.URL.Path == "/native/start":
		id := r.URL.Query().Get("id")
		a.mu.Lock()
		ok := a.find(id) != nil
		a.mu.Unlock()
		if !ok {
			http.NotFound(w, r)
			return
		}
		http.SetCookie(w, &http.Cookie{Name: a.pendingName(), Value: id, Path: "/", MaxAge: 300, Secure: strings.HasPrefix(a.origin, "https:"), HttpOnly: true, SameSite: http.SameSiteLaxMode})
		http.Redirect(w, r, "/", 303)
	case r.URL.Path == "/native/authorize" && (r.Method == "GET" || r.Method == "POST"):
		a.require(http.HandlerFunc(a.authorize)).ServeHTTP(w, r)
	default:
		http.NotFound(w, r)
	}
}
func decode(w http.ResponseWriter, r *http.Request, out any) bool {
	if r.Header.Get("Content-Type") != "application/json" {
		http.Error(w, "Expected JSON", 415)
		return false
	}
	r.Body = http.MaxBytesReader(w, r.Body, 2048)
	d := json.NewDecoder(r.Body)
	if d.Decode(out) != nil || d.Decode(new(any)) != io.EOF {
		http.Error(w, "Invalid request", 400)
		return false
	}
	return true
}
func jsonReply(w http.ResponseWriter, value any) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(value)
}
func (a *Auth) begin(w http.ResponseWriter, r *http.Request) {
	var in struct {
		Challenge string `json:"challenge"`
		Callback  string `json:"callback"`
	}
	if !decode(w, r, &in) {
		return
	}
	if !valid(in.Challenge) || (in.Callback != "" && in.Callback != a.callbackURL()) {
		http.Error(w, "Invalid proof", 400)
		return
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	a.find("")
	if len(a.pending) >= 256 {
		http.Error(w, "Try again later", 429)
		return
	}
	id := token()
	a.pending[id] = &request{challenge: in.Challenge, callback: in.Callback, expires: time.Now().Add(5 * time.Minute)}
	jsonReply(w, map[string]string{"id": id})
}
func (a *Auth) authorize(w http.ResponseWriter, r *http.Request) {
	id := r.URL.Query().Get("id")
	a.mu.Lock()
	p := a.find(id)
	ok := p != nil && p.cookie == ""
	a.mu.Unlock()
	if !ok {
		http.NotFound(w, r)
		return
	}
	done := false
	callback := ""
	if r.Method == "POST" {
		c, e := r.Cookie(a.sessions.CookieName())
		if e != nil {
			http.Error(w, "Sign in again", 401)
			return
		}
		a.mu.Lock()
		p = a.find(id)
		if p == nil || p.cookie != "" {
			a.mu.Unlock()
			http.NotFound(w, r)
			return
		}
		p.cookie = c.Value
		if p.callback != "" {
			callback = p.callback + "?state=" + id
		}
		a.mu.Unlock()
		done = true
	}
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	_ = approval.Execute(w, struct {
		Name, Code, ID string
		Done           bool
		Callback       template.URL
	}{a.name, code(id), id, done, template.URL(callback)})
}
func (a *Auth) exchange(w http.ResponseWriter, r *http.Request) {
	var in struct {
		ID       string `json:"id"`
		Verifier string `json:"verifier"`
	}
	if !decode(w, r, &in) {
		return
	}
	if !valid(in.ID) || !valid(in.Verifier) {
		http.Error(w, "Invalid proof", 400)
		return
	}
	hash := sha256.Sum256([]byte(in.Verifier))
	challenge := base64.RawURLEncoding.EncodeToString(hash[:])
	a.mu.Lock()
	p := a.find(in.ID)
	if p == nil || subtle.ConstantTimeCompare([]byte(p.challenge), []byte(challenge)) != 1 {
		a.mu.Unlock()
		http.NotFound(w, r)
		return
	}
	if p.cookie == "" {
		a.mu.Unlock()
		w.WriteHeader(202)
		return
	}
	cookie := p.cookie
	delete(a.pending, in.ID)
	a.mu.Unlock()
	// Authenticate the approved cookie from scratch. The outer router may have
	// attached a different browser session to the exchange request's context.
	ctx, cancel := context.WithCancel(context.Background())
	stop := context.AfterFunc(r.Context(), cancel)
	defer func() { stop(); cancel() }()
	check := r.Clone(ctx)
	check.Header = make(http.Header)
	check.AddCookie(&http.Cookie{Name: a.sessions.CookieName(), Value: cookie})
	allowed := false
	a.require(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) { allowed = true })).ServeHTTP(discardResponse{}, check)
	session, ok := a.sessions.From(check)
	if !allowed || !ok {
		http.Error(w, "Sign in again", 401)
		return
	}
	value, err := a.sessions.Mint(session)
	if err != nil {
		http.Error(w, "Cannot sign in", 500)
		return
	}
	jsonReply(w, map[string]string{"name": a.sessions.CookieName(), "value": value})
}

type discardResponse struct{}

func (discardResponse) Header() http.Header         { return make(http.Header) }
func (discardResponse) Write(b []byte) (int, error) { return len(b), nil }
func (discardResponse) WriteHeader(int)             {}

var approval = template.Must(template.New("approval").Parse(`<!doctype html><html lang="en"><meta charset="utf-8"><meta name="viewport" content="width=device-width"><title>Sign in to {{.Name}}</title>{{if .Callback}}<meta http-equiv="refresh" content="0;url={{.Callback}}">{{end}}<style>:root{color-scheme:light dark}body{font:17px/1.5 system-ui,sans-serif;margin:0;background:light-dark(#f5f4f8,#19181d);color:light-dark(#24212c,#f7f4fa)}main{max-width:28rem;margin:12vh auto;padding:2rem;border-radius:1.5rem;background:light-dark(white,#28252e);box-shadow:0 8px 40px #0001}h1{font-size:1.7rem;line-height:1.2}h2{font:600 2rem ui-monospace,monospace;letter-spacing:.1em}button{font:600 1rem system-ui;border:0;border-radius:.8rem;padding:.9rem 1.2rem;background:#7052a2;color:white;cursor:pointer}button:focus-visible{outline:3px solid #ad8ddc;outline-offset:4px}@media(max-width:36rem){main{margin:2rem 1rem}}</style><body><main><h1>{{if .Done}}{{.Name}} is signed in{{else}}Sign in to {{.Name}}{{end}}</h1>{{if .Done}}{{if .Callback}}<p><a href="{{.Callback}}">Return to {{.Name}}</a></p>{{else}}<p>Return to the app to continue.</p>{{end}}{{else}}<p>Check that the app shows this same code before continuing.</p><h2>{{.Code}}</h2><form method="post" action="/native/authorize?id={{.ID}}"><button>Sign in to {{.Name}} app</button></form>{{end}}</main></body></html>`))

// A fixed app-specific callback prevents this endpoint becoming an open redirect.
// The callback only wakes the app: session exchange still requires its verifier.
func (a *Auth) callbackURL() string {
	switch a.name {
	case "Docs", "Sheets", "Calendar", "Memoria", "Meet":
		return "works.oficina." + strings.ToLower(a.name) + "://signin"
	default:
		return ""
	}
}