package auth

import (
	"context"
	"crypto"
	"crypto/hmac"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/big"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

type Config struct{ BaseURL, Issuer, ClientID, ClientSecret, SessionSecret string }
type Provider struct{ AuthURL, TokenURL, JWKSURL string }
type Claims struct {
	Sub, Email, Name string
	Exp              int64
	Iss, Aud         string
}

type Manager struct {
	Cfg    Config
	P      Provider
	DB     *sql.DB
	Client *http.Client
}

type contextKey string

const UserIDKey contextKey = "user_id"

func UserID(r *http.Request) string {
	if v, ok := r.Context().Value(UserIDKey).(string); ok {
		return v
	}
	return ""
}

func firstEnv(keys ...string) string {
	for _, k := range keys {
		if v := os.Getenv(k); v != "" {
			return v
		}
	}
	return ""
}
func New(ctx context.Context, db *sql.DB) (*Manager, error) {
	c := Config{firstEnv("APP_BASE_URL", "PUBLIC_URL"), firstEnv("OIDC_ISSUER", "POCKET_ID_ISSUER"), firstEnv("OIDC_CLIENT_ID", "POCKET_ID_CLIENT_ID"), firstEnv("OIDC_CLIENT_SECRET", "POCKET_ID_CLIENT_SECRET"), os.Getenv("SESSION_SECRET")}
	if c.SessionSecret == "" {
		c.SessionSecret = "dev-secret-change-me"
	}
	m := &Manager{Cfg: c, DB: db, Client: &http.Client{Timeout: 10 * time.Second}}
	if c.Issuer != "" {
		if err := m.discover(ctx); err != nil {
			return nil, err
		}
	}
	return m, nil
}
func (m *Manager) discover(ctx context.Context) error {
	req, _ := http.NewRequestWithContext(ctx, "GET", strings.TrimRight(m.Cfg.Issuer, "/")+"/.well-known/openid-configuration", nil)
	resp, err := m.Client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	var v struct {
		AuthorizationEndpoint string `json:"authorization_endpoint"`
		TokenEndpoint         string `json:"token_endpoint"`
		JWKSURI               string `json:"jwks_uri"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&v); err != nil {
		return err
	}
	m.P = Provider{v.AuthorizationEndpoint, v.TokenEndpoint, v.JWKSURI}
	return nil
}
func Rand(n int) string {
	b := make([]byte, n)
	rand.Read(b)
	return base64.RawURLEncoding.EncodeToString(b)
}
func sign(secret, s string) string {
	h := hmac.New(sha256.New, []byte(secret))
	h.Write([]byte(s))
	return base64.RawURLEncoding.EncodeToString(h.Sum(nil))
}
func (m *Manager) SetSignedCookie(w http.ResponseWriter, name, value string, maxAge int) {
	v := value + "." + sign(m.Cfg.SessionSecret, value)
	http.SetCookie(w, &http.Cookie{Name: name, Value: v, Path: "/", HttpOnly: true, SameSite: http.SameSiteLaxMode, Secure: strings.HasPrefix(m.Cfg.BaseURL, "https://"), MaxAge: maxAge})
}
func (m *Manager) ReadSignedCookie(r *http.Request, name string) (string, bool) {
	c, err := r.Cookie(name)
	if err != nil {
		return "", false
	}
	p := strings.LastIndex(c.Value, ".")
	if p < 1 {
		return "", false
	}
	val, sig := c.Value[:p], c.Value[p+1:]
	return val, hmac.Equal([]byte(sig), []byte(sign(m.Cfg.SessionSecret, val)))
}
func (m *Manager) Login(w http.ResponseWriter, r *http.Request) {
	st := Rand(24)
	m.SetSignedCookie(w, "oidc_state", st, 300)
	q := url.Values{"client_id": {m.Cfg.ClientID}, "redirect_uri": {m.Cfg.BaseURL + "/auth/callback"}, "response_type": {"code"}, "scope": {"openid profile email"}, "state": {st}}
	http.Redirect(w, r, m.P.AuthURL+"?"+q.Encode(), 302)
}
func (m *Manager) Callback(w http.ResponseWriter, r *http.Request) {
	st, ok := m.ReadSignedCookie(r, "oidc_state")
	if !ok || st != r.URL.Query().Get("state") {
		http.Error(w, "bad state", 400)
		return
	}
	tok, err := m.exchange(r.Context(), r.URL.Query().Get("code"))
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	cl, err := m.verify(r.Context(), tok)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	res, err := m.DB.ExecContext(r.Context(), `INSERT INTO users(subject,email,name) VALUES(?,?,?) ON DUPLICATE KEY UPDATE email=VALUES(email), name=VALUES(name), id=LAST_INSERT_ID(id)`, cl.Sub, cl.Email, cl.Name)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	uid, err := res.LastInsertId()
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	m.SetSignedCookie(w, "session", fmt.Sprint(uid), 86400*30)
	http.Redirect(w, r, "/", 302)
}
func (m *Manager) exchange(ctx context.Context, code string) (string, error) {
	data := url.Values{"grant_type": {"authorization_code"}, "code": {code}, "redirect_uri": {m.Cfg.BaseURL + "/auth/callback"}, "client_id": {m.Cfg.ClientID}, "client_secret": {m.Cfg.ClientSecret}}
	req, _ := http.NewRequestWithContext(ctx, "POST", m.P.TokenURL, strings.NewReader(data.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	resp, err := m.Client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode > 299 {
		return "", fmt.Errorf("token error: %s", body)
	}
	var v struct {
		IDToken string `json:"id_token"`
	}
	json.Unmarshal(body, &v)
	return v.IDToken, nil
}
func (m *Manager) verify(ctx context.Context, jwt string) (Claims, error) {
	parts := strings.Split(jwt, ".")
	if len(parts) != 3 {
		return Claims{}, errors.New("bad jwt")
	}
	headB, _ := base64.RawURLEncoding.DecodeString(parts[0])
	var h struct{ Kid, Alg string }
	json.Unmarshal(headB, &h)
	if h.Alg != "RS256" {
		return Claims{}, errors.New("only RS256 supported")
	}
	key, err := m.jwk(ctx, h.Kid)
	if err != nil {
		return Claims{}, err
	}
	sig, _ := base64.RawURLEncoding.DecodeString(parts[2])
	sum := sha256.Sum256([]byte(parts[0] + "." + parts[1]))
	if err := rsa.VerifyPKCS1v15(key, crypto.SHA256, sum[:], sig); err != nil {
		return Claims{}, err
	}
	pay, _ := base64.RawURLEncoding.DecodeString(parts[1])
	var raw map[string]any
	json.Unmarshal(pay, &raw)
	c := Claims{Sub: fmt.Sprint(raw["sub"]), Email: fmt.Sprint(raw["email"]), Name: fmt.Sprint(raw["name"]), Iss: fmt.Sprint(raw["iss"]), Aud: fmt.Sprint(raw["aud"])}
	if exp, ok := raw["exp"].(float64); ok {
		c.Exp = int64(exp)
	}
	if c.Exp < time.Now().Unix() {
		return c, errors.New("token expired")
	}
	return c, nil
}
func (m *Manager) jwk(ctx context.Context, kid string) (*rsa.PublicKey, error) {
	req, _ := http.NewRequestWithContext(ctx, "GET", m.P.JWKSURL, nil)
	resp, err := m.Client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	var jwks struct {
		Keys []map[string]string `json:"keys"`
	}
	json.NewDecoder(resp.Body).Decode(&jwks)
	for _, k := range jwks.Keys {
		if k["kid"] == kid {
			nB, _ := base64.RawURLEncoding.DecodeString(k["n"])
			eB, _ := base64.RawURLEncoding.DecodeString(k["e"])
			e := 0
			for _, b := range eB {
				e = e*256 + int(b)
			}
			return &rsa.PublicKey{N: new(big.Int).SetBytes(nB), E: e}, nil
		}
	}
	return nil, errors.New("jwk not found")
}
func (m *Manager) Require(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		uid, ok := m.ReadSignedCookie(r, "session")
		if !ok || uid == "" {
			http.Redirect(w, r, "/login", 302)
			return
		}
		ctx := context.WithValue(r.Context(), UserIDKey, uid)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}
func Logout(w http.ResponseWriter, r *http.Request) {
	http.SetCookie(w, &http.Cookie{Name: "session", Path: "/", MaxAge: -1})
	http.Redirect(w, r, "/login", 302)
}

func (m *Manager) RequireAdmin(next http.Handler) http.Handler {
	return m.Require(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		uid, _ := m.ReadSignedCookie(r, "session")
		var email string
		m.DB.QueryRowContext(r.Context(), `select email from users where id=?`, uid).Scan(&email)
		var allowed string
		m.DB.QueryRowContext(r.Context(), `select setting_value from app_settings where setting_key='admin_emails'`).Scan(&allowed)
		if strings.TrimSpace(allowed) != "" {
			ok := false
			for _, part := range strings.Split(allowed, ",") {
				if strings.EqualFold(strings.TrimSpace(part), strings.TrimSpace(email)) {
					ok = true
					break
				}
			}
			if !ok {
				http.Error(w, "admin access required", http.StatusForbidden)
				return
			}
		}
		next.ServeHTTP(w, r)
	}))
}