package oauth

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"sync"
	"time"

	redispkg "github.com/ClareAI/astra-insight-mcp/internal/pkg/redis"
)

var frontendHosts = map[string]string{
	"prod":  "https://astra.wati.io",
	"stage": "https://astra.wati.io",
	"dev":   "https://dev-astra.watiapp.io",
	"qa":    "https://qa-astra.watiapp.io",
}

var gatewayBases = map[string]string{
	"prod":  "https://astra-gateway.wati.io",
	"stage": "https://astra-gateway.wati.io",
	"dev":   "https://astra-gateway.watiapp.io",
	"qa":    "https://qa-astra-gateway.watiapp.io",
}

const (
	DefaultClientID   = "astra_cli_public"
	DefaultScope      = "all"
	DefaultListenHost = "127.0.0.1"
	DefaultListenPort = 8787
	DefaultTimeout    = 120
	callbackPath      = "/callback"
)

type Config struct {
	Env        string
	ClientID   string
	Scope      string
	ListenHost string
	ListenPort int
	Timeout    int // seconds
	NoOpen     bool
}

func (c *Config) applyDefaults() {
	if c.Env == "" {
		c.Env = "dev"
	}
	if c.ClientID == "" {
		c.ClientID = DefaultClientID
	}
	if c.Scope == "" {
		c.Scope = DefaultScope
	}
	if c.ListenHost == "" {
		c.ListenHost = DefaultListenHost
	}
	if c.ListenPort == 0 {
		c.ListenPort = DefaultListenPort
	}
	if c.Timeout == 0 {
		c.Timeout = DefaultTimeout
	}
}

type TokenData struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token,omitempty"`
	TokenType    string `json:"token_type,omitempty"`
	Scope        string `json:"scope,omitempty"`
	ExpiresIn    int    `json:"expires_in,omitempty"`
	ClientID     string `json:"client_id,omitempty"`
	Env          string `json:"env,omitempty"`
	ObtainedAt   int64  `json:"obtained_at,omitempty"`
}

func TokenFilePath() string {
	home, err := os.UserHomeDir()
	if err != nil {
		return ""
	}
	return filepath.Join(home, ".astra", "oauth_token.json")
}

func LoadStoredToken() (*TokenData, error) {
	path := TokenFilePath()
	if path == "" {
		return nil, fmt.Errorf("cannot determine home directory")
	}
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var tok TokenData
	if err := json.Unmarshal(data, &tok); err != nil {
		return nil, err
	}
	return &tok, nil
}

func SaveToken(tok *TokenData) error {
	path := TokenFilePath()
	if path == "" {
		return fmt.Errorf("cannot determine home directory")
	}
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return err
	}
	if tok.ObtainedAt == 0 {
		tok.ObtainedAt = time.Now().Unix()
	}
	data, err := json.MarshalIndent(tok, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(path, data, 0600)
}

func generatePKCE() (verifier, challenge string, err error) {
	buf := make([]byte, 64)
	if _, err = rand.Read(buf); err != nil {
		return "", "", fmt.Errorf("crypto/rand: %w", err)
	}
	verifier = base64.RawURLEncoding.EncodeToString(buf)
	h := sha256.Sum256([]byte(verifier))
	challenge = base64.RawURLEncoding.EncodeToString(h[:])
	return verifier, challenge, nil
}

func generateState() (string, error) {
	buf := make([]byte, 16)
	if _, err := rand.Read(buf); err != nil {
		return "", fmt.Errorf("crypto/rand: %w", err)
	}
	return base64.RawURLEncoding.EncodeToString(buf), nil
}

// RunBrowserAuth performs the full OAuth PKCE flow:
// 1. Start local callback server
// 2. Open browser to authorize URL
// 3. Wait for callback with code+state
// 4. Exchange code for token
// 5. Save token to disk
func RunBrowserAuth(cfg Config) error {
	cfg.applyDefaults()

	gateway, ok := gatewayBases[cfg.Env]
	if !ok {
		return fmt.Errorf("unknown env %q (valid: prod, stage, dev, qa)", cfg.Env)
	}

	authCtx, err := BuildAuthorizeContext(cfg)
	if err != nil {
		return err
	}

	fmt.Fprintf(os.Stderr, "[auth] env          : %s\n", cfg.Env)
	fmt.Fprintf(os.Stderr, "[auth] client_id    : %s\n", cfg.ClientID)
	fmt.Fprintf(os.Stderr, "[auth] redirect_uri : %s\n", authCtx.RedirectURI)
	fmt.Fprintf(os.Stderr, "[auth] scope        : %s\n", cfg.Scope)
	fmt.Fprintf(os.Stderr, "[auth] PKCE         : S256\n")
	fmt.Fprintf(os.Stderr, "[auth] timeout      : %ds\n", cfg.Timeout)
	fmt.Fprintf(os.Stderr, "\n[auth] Authorize URL:\n  %s\n\n", authCtx.URL)

	if cfg.NoOpen {
		fmt.Fprintf(os.Stderr, "[auth] --no-open set. Open the URL above in your browser manually.\n")
	} else {
		fmt.Fprintf(os.Stderr, "[auth] Opening browser...\n")
		if err := openBrowser(authCtx.URL); err != nil {
			fmt.Fprintf(os.Stderr, "[auth] Failed to open browser: %v\n", err)
			fmt.Fprintf(os.Stderr, "[auth] Please open the URL above manually.\n")
		}
	}

	// Start callback server and wait
	fmt.Fprintf(os.Stderr, "[auth] Listening on %s:%d%s (timeout %ds)...\n",
		cfg.ListenHost, cfg.ListenPort, callbackPath, cfg.Timeout)

	cbParams, err := waitForCallback(cfg.ListenHost, cfg.ListenPort, cfg.Timeout)
	if err != nil {
		return fmt.Errorf("callback failed: %w", err)
	}
	if cbParams == nil {
		fmt.Fprintf(os.Stderr, "\n[auth] TIMEOUT after %ds. No callback received.\n", cfg.Timeout)
		return fmt.Errorf("timeout waiting for OAuth callback")
	}

	if errMsg := cbParams["error"]; errMsg != "" {
		desc := cbParams["error_description"]
		return fmt.Errorf("OAuth error: %s (%s)", errMsg, desc)
	}

	code := cbParams["code"]
	returnedState := cbParams["state"]
	if code == "" {
		return fmt.Errorf("no 'code' in callback parameters")
	}
	if returnedState != authCtx.State {
		return fmt.Errorf("state mismatch (expected=%s, got=%s)", authCtx.State, returnedState)
	}

	truncCode := code
	if len(truncCode) > 20 {
		truncCode = truncCode[:20] + "..."
	}
	fmt.Fprintf(os.Stderr, "[auth] Callback received. code=%s state=OK\n", truncCode)

	// Exchange code for token
	fmt.Fprintf(os.Stderr, "[auth] Exchanging code for token...\n")
	tok, err := exchangeCode(gateway, code, cfg.ClientID, authCtx.RedirectURI, authCtx.CodeVerifier)
	if err != nil {
		return fmt.Errorf("token exchange failed: %w", err)
	}

	truncAT := tok.AccessToken
	if len(truncAT) > 8 {
		truncAT = truncAT[:8] + "..."
	}
	fmt.Fprintf(os.Stderr, "[auth] Token exchange SUCCESS.\n")
	fmt.Fprintf(os.Stderr, "[auth]   access_token : %s\n", truncAT)
	fmt.Fprintf(os.Stderr, "[auth]   expires_in   : %ds\n", tok.ExpiresIn)

	tok.ClientID = cfg.ClientID
	tok.Env = cfg.Env
	tok.ObtainedAt = time.Now().Unix()
	if err := SaveToken(tok); err != nil {
		return fmt.Errorf("failed to save token: %w", err)
	}
	fmt.Fprintf(os.Stderr, "[auth] Token saved to %s\n", TokenFilePath())
	fmt.Fprintf(os.Stderr, "\n[auth] DONE. You can now use the CLI without --api-key.\n")
	return nil
}

// AuthorizeContext holds all PKCE and state material generated for an
// authorization request. Callers that need to complete the flow (callback
// validation + token exchange) must retain CodeVerifier and State.
type AuthorizeContext struct {
	URL           string // full authorize URL to open in a browser
	State         string // opaque state for CSRF validation
	CodeVerifier  string // PKCE code_verifier (needed for token exchange)
	CodeChallenge string // PKCE code_challenge sent in the URL
	RedirectURI   string // callback URI the authorization server will redirect to
}

// BuildAuthorizeContext generates PKCE credentials, state, and the full
// authorize URL. Use this when you need to keep the verifier/state for a
// subsequent token exchange or callback validation.
func BuildAuthorizeContext(cfg Config) (*AuthorizeContext, error) {
	cfg.applyDefaults()
	frontend, ok := frontendHosts[cfg.Env]
	if !ok {
		return nil, fmt.Errorf("unknown env %q", cfg.Env)
	}
	verifier, challenge, err := generatePKCE()
	if err != nil {
		return nil, fmt.Errorf("generate PKCE: %w", err)
	}
	state, err := generateState()
	if err != nil {
		return nil, fmt.Errorf("generate state: %w", err)
	}
	redirectURI := fmt.Sprintf("http://%s:%d%s", cfg.ListenHost, cfg.ListenPort, callbackPath)
	params := url.Values{
		"client_id":             {cfg.ClientID},
		"redirect_uri":          {redirectURI},
		"scope":                 {cfg.Scope},
		"state":                 {state},
		"code_challenge":        {challenge},
		"code_challenge_method": {"S256"},
		"response_type":         {"code"},
	}
	return &AuthorizeContext{
		URL:           fmt.Sprintf("%s/oauth/authorize?%s", frontend, params.Encode()),
		State:         state,
		CodeVerifier:  verifier,
		CodeChallenge: challenge,
		RedirectURI:   redirectURI,
	}, nil
}

// BuildAuthorizeURL returns only the authorize URL for display or testing.
// The generated PKCE verifier and state are discarded — use
// BuildAuthorizeContext instead when you need to complete the full flow.
func BuildAuthorizeURL(cfg Config) (string, error) {
	ctx, err := BuildAuthorizeContext(cfg)
	if err != nil {
		return "", err
	}
	return ctx.URL, nil
}

// --- callback server ---

type callbackResult struct {
	params map[string]string
}

func waitForCallback(host string, port, timeoutSec int) (map[string]string, error) {
	addr := fmt.Sprintf("%s:%d", host, port)
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, fmt.Errorf("listen %s: %w", addr, err)
	}

	var (
		result    *callbackResult
		mu        sync.Mutex
		done      = make(chan struct{})
		closeOnce sync.Once
	)

	mux := http.NewServeMux()
	mux.HandleFunc(callbackPath, func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		defer mu.Unlock()

		params := make(map[string]string)
		for k, v := range r.URL.Query() {
			if len(v) > 0 {
				params[k] = v[0]
			}
		}
		result = &callbackResult{params: params}

		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, `<html><body style="font-family:system-ui,sans-serif;display:flex;`+
			`justify-content:center;align-items:center;height:90vh">`+
			`<div style="text-align:center">`+
			`<h1>Authorization Successful</h1>`+
			`<p>You can close this tab and return to the terminal.</p>`+
			`</div></body></html>`)

		go func() {
			time.Sleep(200 * time.Millisecond)
			closeOnce.Do(func() { close(done) })
		}()
	})

	srv := &http.Server{Handler: mux}

	go func() {
		_ = srv.Serve(listener)
	}()

	select {
	case <-done:
		_ = srv.Close()
		mu.Lock()
		defer mu.Unlock()
		if result != nil {
			return result.params, nil
		}
		return nil, nil
	case <-time.After(time.Duration(timeoutSec) * time.Second):
		_ = srv.Close()
		return nil, nil
	}
}

// RefreshAccessToken uses a refresh_token to obtain a new access_token and
// refresh_token pair. The returned TokenData is already saved to disk.
func RefreshAccessToken(tok *TokenData) (*TokenData, error) {
	if tok.RefreshToken == "" {
		return nil, fmt.Errorf("no refresh_token available")
	}
	clientID := tok.ClientID
	if clientID == "" {
		clientID = DefaultClientID
	}
	env := tok.Env
	if env == "" {
		env = "dev"
	}
	gateway, ok := gatewayBases[env]
	if !ok {
		return nil, fmt.Errorf("unknown env %q", env)
	}

	reqBody := map[string]string{
		"grant_type":    "refresh_token",
		"refresh_token": tok.RefreshToken,
		"client_id":     clientID,
	}
	jsonBytes, err := json.Marshal(reqBody)
	if err != nil {
		return nil, fmt.Errorf("marshal refresh request: %w", err)
	}

	tokenURL := gateway + "/oauth/token"
	client := &http.Client{Timeout: 30 * time.Second}
	req, err := http.NewRequest(http.MethodPost, tokenURL, bytes.NewReader(jsonBytes))
	if err != nil {
		return nil, fmt.Errorf("build refresh request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("POST %s: %w", tokenURL, err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("read refresh response: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("refresh endpoint returned %d: %s", resp.StatusCode, string(respBody))
	}

	var newTok TokenData
	if err := json.Unmarshal(respBody, &newTok); err != nil {
		return nil, fmt.Errorf("decode refresh response: %w", err)
	}
	if newTok.AccessToken == "" {
		return nil, fmt.Errorf("no access_token in refresh response")
	}

	newTok.ClientID = clientID
	newTok.Env = env
	newTok.ObtainedAt = time.Now().Unix()
	if err := SaveToken(&newTok); err != nil {
		return nil, fmt.Errorf("save refreshed token: %w", err)
	}
	return &newTok, nil
}

// IsExpired returns true if the access token has likely expired based on
// obtained_at and expires_in. Uses a 5-minute buffer.
func (t *TokenData) IsExpired() bool {
	if t.ObtainedAt == 0 || t.ExpiresIn == 0 {
		return false
	}
	expiresAt := time.Unix(t.ObtainedAt, 0).Add(time.Duration(t.ExpiresIn) * time.Second)
	return time.Now().After(expiresAt.Add(-5 * time.Minute))
}

// EnsureValidToken loads the stored token and refreshes it if expired.
// Returns the valid access token string. If no token file exists or
// refresh fails, returns empty string.
func EnsureValidToken() string {
	tok, err := LoadStoredToken()
	if err != nil || tok.AccessToken == "" {
		return ""
	}
	if !tok.IsExpired() {
		return tok.AccessToken
	}
	if tok.RefreshToken == "" {
		return tok.AccessToken
	}
	newTok, err := RefreshAccessToken(tok)
	if err != nil {
		fmt.Fprintf(os.Stderr, "[oauth] token refresh failed: %v\n", err)
		return tok.AccessToken
	}
	fmt.Fprintf(os.Stderr, "[oauth] token refreshed successfully\n")
	return newTok.AccessToken
}

// ---- Server-side OAuth (for webhook / WhatsApp users) ----

// PendingAuth holds PKCE material for an in-flight OAuth flow keyed by state.
type PendingAuth struct {
	WaID         string
	State        string
	CodeVerifier string
	RedirectURI  string
	Env          string
	ClientID     string
	CreatedAt    int64
}

// PendingAuthStore persists in-flight OAuth PKCE state in Redis.
// Keys auto-expire via TTL so no manual cleanup is needed.
type PendingAuthStore struct {
	redis *redispkg.Client
}

const pendingAuthKeyPrefix = "astra:insight_mcp:pending_auth:"
const pendingAuthTTL = 5 * time.Minute

func NewPendingAuthStore(rdb *redispkg.Client) *PendingAuthStore {
	return &PendingAuthStore{redis: rdb}
}

func (s *PendingAuthStore) Put(p *PendingAuth) {
	data, err := json.Marshal(p)
	if err != nil {
		slog.Error("pending_auth: marshal failed", "state", p.State, "error", err)
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()

	if err := s.redis.Set(ctx, pendingAuthKeyPrefix+p.State, string(data), pendingAuthTTL); err != nil {
		slog.Error("pending_auth: redis set failed", "state", p.State, "error", err)
	}
}

func (s *PendingAuthStore) Pop(state string) *PendingAuth {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()

	val, err := s.redis.GetDel(ctx, pendingAuthKeyPrefix+state)
	if err != nil {
		slog.Error("pending_auth: redis getdel failed", "state", state, "error", err)
		return nil
	}
	if val == "" {
		return nil
	}
	var p PendingAuth
	if err := json.Unmarshal([]byte(val), &p); err != nil {
		slog.Error("pending_auth: unmarshal failed", "state", state, "error", err)
		return nil
	}
	return &p
}

// Cleanup is a no-op when using Redis — TTL handles expiry automatically.
func (s *PendingAuthStore) Cleanup(maxAgeSec int64) {}

// BuildServerAuthorizeContext generates an authorize URL that redirects to
// a server-hosted callback (not localhost). Used for WhatsApp OAuth flows.
func BuildServerAuthorizeContext(cfg Config, callbackBaseURL string) (*AuthorizeContext, error) {
	cfg.applyDefaults()
	frontend, ok := frontendHosts[cfg.Env]
	if !ok {
		return nil, fmt.Errorf("unknown env %q", cfg.Env)
	}
	verifier, challenge, err := generatePKCE()
	if err != nil {
		return nil, fmt.Errorf("generate PKCE: %w", err)
	}
	state, err := generateState()
	if err != nil {
		return nil, fmt.Errorf("generate state: %w", err)
	}
	redirectURI := callbackBaseURL + "/oauth/callback"
	params := url.Values{
		"client_id":             {cfg.ClientID},
		"redirect_uri":          {redirectURI},
		"scope":                 {cfg.Scope},
		"state":                 {state},
		"code_challenge":        {challenge},
		"code_challenge_method": {"S256"},
		"response_type":         {"code"},
	}
	return &AuthorizeContext{
		URL:           fmt.Sprintf("%s/oauth/authorize?%s", frontend, params.Encode()),
		State:         state,
		CodeVerifier:  verifier,
		CodeChallenge: challenge,
		RedirectURI:   redirectURI,
	}, nil
}

// ExchangeCodeForToken exchanges an authorization code for a token using
// the given gateway base URL. Exported for use by the server callback handler.
func ExchangeCodeForToken(env, code, clientID, redirectURI, codeVerifier string) (*TokenData, error) {
	gateway, ok := gatewayBases[env]
	if !ok {
		return nil, fmt.Errorf("unknown env %q", env)
	}
	return exchangeCode(gateway, code, clientID, redirectURI, codeVerifier)
}

// --- token exchange ---

func exchangeCode(gatewayBase, code, clientID, redirectURI, codeVerifier string) (*TokenData, error) {
	reqBody := map[string]string{
		"grant_type":    "authorization_code",
		"code":          code,
		"client_id":     clientID,
		"redirect_uri":  redirectURI,
		"code_verifier": codeVerifier,
	}
	jsonBytes, err := json.Marshal(reqBody)
	if err != nil {
		return nil, fmt.Errorf("marshal request: %w", err)
	}

	tokenURL := gatewayBase + "/oauth/token"
	client := &http.Client{Timeout: 30 * time.Second}
	req, err := http.NewRequest(http.MethodPost, tokenURL, bytes.NewReader(jsonBytes))
	if err != nil {
		return nil, fmt.Errorf("build request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("POST %s: %w", tokenURL, err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("read response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("token endpoint returned %d: %s", resp.StatusCode, string(respBody))
	}

	var tok TokenData
	if err := json.Unmarshal(respBody, &tok); err != nil {
		return nil, fmt.Errorf("decode token response: %w", err)
	}
	if tok.AccessToken == "" {
		return nil, fmt.Errorf("no access_token in response: %s", string(respBody))
	}
	return &tok, nil
}
