// Package ws provides bounded, managed WebSocket connections.
package ws

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
	"unicode/utf8"

	"github.com/gorilla/websocket"
)

type MessageType int

const (
	Text   MessageType = websocket.TextMessage
	Binary MessageType = websocket.BinaryMessage
)

type Message struct {
	Type MessageType
	Data []byte
}
type Handler func(context.Context, *Session, Message) error

// Handlers execute serially for each session. OnClose runs after all I/O and
// message callbacks have stopped; it must not call Wait on its own session.
type Handlers struct {
	OnOpen    func(context.Context, *Session) error
	OnMessage Handler
	OnClose   func(*Session, error)
}

type SessionStats struct {
	MessagesSent, MessagesReceived, BytesSent, BytesReceived uint64
	QueuedMessages                                           int
	QueuedBytes                                              int64
	ReadQueuedMessages                                       int
	ReadQueuedBytes                                          int64
}

type outgoing struct {
	message Message
	ctx     context.Context
	result  chan error
}
type controlFrame struct {
	kind int
	data []byte
}

var nextSessionID atomic.Uint64

// Session owns exactly one reader and one writer. Sends are ordered by queue
// admission. Message callbacks never run on the reader goroutine.
type Session struct {
	conn                                            *websocket.Conn
	cfg                                             SessionConfig
	handlers                                        Handlers
	id                                              string
	parent                                          context.Context
	ctx                                             context.Context
	cancel                                          context.CancelFunc
	done, stopped, closingCh, peerClosed, closeSent chan struct{}
	startOnce, stopOnce, peerOnce                   sync.Once
	wg                                              sync.WaitGroup
	mu                                              sync.Mutex
	closing                                         bool
	closeCode                                       int
	closeReason                                     string
	err                                             error
	stats                                           SessionStats
	out                                             chan outgoing
	in                                              chan Message
	controls                                        chan controlFrame
}

// newSession requires a validated config. Ownership transfers before start so
// callers can register the session without racing completion.
func newSession(conn *websocket.Conn, cfg SessionConfig, handlers Handlers, parent context.Context) *Session {
	if parent == nil {
		parent = context.Background()
	}
	ctx, cancel := context.WithCancel(parent)
	cfg = cfg.defaults()
	return &Session{conn: conn, cfg: cfg, handlers: handlers, id: "ws-" + strconv.FormatUint(nextSessionID.Add(1), 10),
		parent: parent, ctx: ctx, cancel: cancel, done: make(chan struct{}), stopped: make(chan struct{}),
		closingCh: make(chan struct{}), peerClosed: make(chan struct{}), closeSent: make(chan struct{}),
		out: make(chan outgoing, cfg.WriteQueueSize), in: make(chan Message, cfg.ReadQueueSize), controls: make(chan controlFrame, 8)}
}

func (s *Session) ID() string               { return s.id }
func (s *Session) Context() context.Context { return s.ctx }
func (s *Session) Done() <-chan struct{}    { return s.done }
func (s *Session) Err() error               { s.mu.Lock(); defer s.mu.Unlock(); return s.err }
func (s *Session) Subprotocol() string      { return s.conn.Subprotocol() }
func (s *Session) Stats() SessionStats      { s.mu.Lock(); defer s.mu.Unlock(); return s.stats }

func (s *Session) start() {
	s.startOnce.Do(func() {
		s.wg.Add(4)
		go func() { defer s.wg.Done(); s.readLoop() }()
		go func() { defer s.wg.Done(); s.writeLoop() }()
		go func() { defer s.wg.Done(); s.dispatchLoop() }()
		go func() { defer s.wg.Done(); s.monitor() }()
		go func() {
			s.wg.Wait()
			for {
				select {
				case item := <-s.out:
					s.complete(item, ErrClosed)
				default:
					goto drained
				}
			}
		drained:
			// Release any discarded inbound payloads before publishing completion.
			for {
				select {
				case msg := <-s.in:
					s.releaseRead(msg)
				default:
					goto notified
				}
			}
		notified:
			if s.handlers.OnClose != nil {
				func() {
					defer func() {
						if p := recover(); p != nil {
							s.recordError(&OpError{Op: "on_close", Cause: fmt.Errorf("callback panic: %v", p)})
						}
					}()
					s.handlers.OnClose(s, s.Err())
				}()
			}
			close(s.done)
		}()
	})
}

// Send waits for the actual write. A cancellation before writing skips the
// message. Cancellation during writing terminates the connection; delivery may
// already have occurred and the caller must not automatically replay it.
func (s *Session) Send(ctx context.Context, msg Message) error {
	if ctx == nil {
		ctx = context.Background()
	}
	item, err := s.enqueue(ctx, msg, true)
	if err != nil {
		return err
	}
	select {
	case err := <-item.result:
		return err
	case <-ctx.Done():
		return ctx.Err()
	}
}

// TrySend only acknowledges queue admission, not delivery. The payload is copied.
func (s *Session) TrySend(msg Message) error {
	_, err := s.enqueue(context.Background(), msg, false)
	return err
}
func (s *Session) SendJSON(ctx context.Context, value any) error {
	data, err := json.Marshal(value)
	if err != nil {
		return &OpError{Op: "encode", Cause: err}
	}
	return s.Send(ctx, Message{Type: Text, Data: data})
}

func (s *Session) enqueue(ctx context.Context, msg Message, wait bool) (outgoing, error) {
	if err := ctx.Err(); err != nil {
		return outgoing{}, err
	}
	if msg.Type != Text && msg.Type != Binary {
		return outgoing{}, ErrInvalidMessage
	}
	if int64(len(msg.Data)) > s.cfg.MaxMessageBytes {
		return outgoing{}, ErrMessageTooLarge
	}
	if msg.Type == Text && !utf8.Valid(msg.Data) {
		return outgoing{}, ErrInvalidMessage
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closing {
		return outgoing{}, ErrClosed
	}
	if s.stats.QueuedMessages >= s.cfg.WriteQueueSize || int64(len(msg.Data)) > s.cfg.WriteQueueBytes-s.stats.QueuedBytes {
		return outgoing{}, ErrQueueFull
	}
	item := outgoing{message: Message{Type: msg.Type, Data: append([]byte(nil), msg.Data...)}, ctx: ctx}
	if wait {
		item.result = make(chan error, 1)
	}
	select {
	case s.out <- item:
		s.stats.QueuedMessages++
		s.stats.QueuedBytes += int64(len(msg.Data))
		return item, nil
	default:
		return outgoing{}, ErrQueueFull
	}
}

func validCloseCode(code int) bool {
	switch code {
	case 1000, 1001, 1002, 1003, 1007, 1008, 1009, 1010, 1011, 1012, 1013, 1014:
		return true
	}
	return code >= 3000 && code <= 4999
}

// Close stops admission and begins a bounded graceful close, draining already
// admitted sends. It is idempotent and never waits for a callback to finish.
func (s *Session) Close(code int, reason string) error {
	if !validCloseCode(code) || len(reason) > 123 || !utf8.ValidString(reason) {
		return ErrInvalidMessage
	}
	s.beginClose(code, reason, nil)
	return nil
}
func (s *Session) beginClose(code int, reason string, cause error) {
	s.mu.Lock()
	if cause != nil && s.err == nil {
		s.err = cause
	}
	if !s.closing {
		s.closing = true
		s.closeCode = code
		s.closeReason = reason
		close(s.closingCh)
	}
	s.mu.Unlock()
	s.cancel()
}
func (s *Session) recordError(err error) {
	s.mu.Lock()
	if s.err == nil {
		s.err = err
	}
	s.mu.Unlock()
}
func (s *Session) terminate(err error) {
	s.beginClose(1001, "", err)
	s.stopOnce.Do(func() { close(s.stopped); _ = s.conn.Close() })
}

// Abort immediately closes the socket; Wait still waits for callbacks to return.
func (s *Session) Abort() { s.terminate(ErrClosed) }
func (s *Session) Wait(ctx context.Context) error {
	if ctx == nil {
		ctx = context.Background()
	}
	select {
	case <-s.done:
		return s.Err()
	default:
	}
	select {
	case <-s.done:
		return s.Err()
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (s *Session) monitor() {
	select {
	case <-s.stopped:
		return
	case <-s.parent.Done():
		s.beginClose(1001, "", s.parent.Err())
	case <-s.closingCh:
	}
	timer := time.NewTimer(s.cfg.CloseTimeout)
	defer timer.Stop()
	select {
	case <-s.stopped:
		return
	case <-timer.C:
		s.terminate(&OpError{Op: "close", Cause: context.DeadlineExceeded})
		return
	case <-s.closeSent:
	}
	select {
	case <-s.stopped:
		return
	case <-timer.C:
		s.terminate(&OpError{Op: "close", Cause: context.DeadlineExceeded})
	case <-s.peerClosed:
		s.terminate(nil)
	}
}

func (s *Session) readLoop() {
	s.conn.SetReadLimit(s.cfg.MaxMessageBytes)
	if !s.cfg.DisableHeartbeat {
		_ = s.conn.SetReadDeadline(time.Now().Add(s.cfg.PongTimeout))
		s.conn.SetPongHandler(func(string) error { return s.conn.SetReadDeadline(time.Now().Add(s.cfg.PongTimeout)) })
	}
	s.conn.SetPingHandler(func(data string) error {
		select {
		case s.controls <- controlFrame{websocket.PongMessage, []byte(data)}:
			return nil
		default:
			return ErrQueueFull
		}
	})
	s.conn.SetCloseHandler(func(code int, _ string) error {
		var cause error
		s.mu.Lock()
		locallyClosing := s.closing
		s.mu.Unlock()
		// A remote going-away frame can signal a restarting service. Preserve
		// its code for Client.Run's retry policy; our own close acknowledgement
		// remains a normal local shutdown.
		if code != 1000 && code != 1005 && (code != 1001 || !locallyClosing) {
			cause = &CloseError{Code: code}
		}
		reply := code
		if !validCloseCode(reply) {
			reply = 1000
		}
		s.beginClose(reply, "", cause)
		s.peerOnce.Do(func() { close(s.peerClosed) })
		return &websocket.CloseError{Code: code}
	})
	for {
		kind, reader, err := s.conn.NextReader()
		if err != nil {
			select {
			case <-s.peerClosed:
				return
			case <-s.stopped:
				return
			default:
			}
			var ce *websocket.CloseError
			if errors.As(err, &ce) {
				s.terminate(&CloseError{Code: ce.Code, Cause: err})
			} else if errors.Is(err, websocket.ErrReadLimit) {
				s.terminate(ErrMessageTooLarge)
			} else {
				s.terminate(&OpError{Op: "read", Cause: err})
			}
			return
		}
		data, err := io.ReadAll(io.LimitReader(reader, s.cfg.MaxMessageBytes+1))
		if err != nil {
			if errors.Is(err, websocket.ErrReadLimit) {
				s.beginClose(1009, "", ErrMessageTooLarge)
			} else {
				s.terminate(&OpError{Op: "read", Cause: err})
			}
			return
		}
		if int64(len(data)) > s.cfg.MaxMessageBytes {
			s.beginClose(1009, "", ErrMessageTooLarge)
			return
		}
		if kind == websocket.TextMessage && !utf8.Valid(data) {
			s.beginClose(1007, "", ErrInvalidMessage)
			return
		}
		s.mu.Lock()
		if s.closing {
			s.mu.Unlock()
			continue
		}
		if s.stats.ReadQueuedMessages >= s.cfg.ReadQueueSize || int64(len(data)) > s.cfg.ReadQueueBytes-s.stats.ReadQueuedBytes {
			s.mu.Unlock()
			s.beginClose(1008, "", ErrQueueFull)
			return
		}
		msg := Message{Type: MessageType(kind), Data: data}
		select {
		case s.in <- msg:
			s.stats.ReadQueuedMessages++
			s.stats.ReadQueuedBytes += int64(len(data))
			s.stats.MessagesReceived++
			s.stats.BytesReceived += uint64(len(data))
			s.mu.Unlock()
		default:
			s.mu.Unlock()
			s.beginClose(1008, "", ErrQueueFull)
			return
		}
	}
}

func (s *Session) releaseRead(msg Message) {
	s.mu.Lock()
	s.stats.ReadQueuedMessages--
	s.stats.ReadQueuedBytes -= int64(len(msg.Data))
	s.mu.Unlock()
}
func (s *Session) callback(fn func() error) (err error) {
	defer func() {
		if p := recover(); p != nil {
			err = fmt.Errorf("callback panic: %v", p)
		}
	}()
	return fn()
}
func (s *Session) dispatchLoop() {
	if s.ctx.Err() != nil {
		return
	}
	if s.handlers.OnOpen != nil {
		if err := s.callback(func() error { return s.handlers.OnOpen(s.ctx, s) }); err != nil {
			s.beginClose(1011, "", &OpError{Op: "on_open", Cause: err})
			return
		}
	}
	for {
		select {
		case <-s.ctx.Done():
			return
		case msg := <-s.in:
			if s.ctx.Err() != nil {
				s.releaseRead(msg)
				return
			}
			var err error
			if s.handlers.OnMessage != nil {
				err = s.callback(func() error { return s.handlers.OnMessage(s.ctx, s, msg) })
			}
			s.releaseRead(msg)
			if err != nil {
				s.beginClose(1011, "", &OpError{Op: "on_message", Cause: err})
				return
			}
		}
	}
}

func (s *Session) complete(item outgoing, err error) {
	s.mu.Lock()
	s.stats.QueuedMessages--
	s.stats.QueuedBytes -= int64(len(item.message.Data))
	if err == nil {
		s.stats.MessagesSent++
		s.stats.BytesSent += uint64(len(item.message.Data))
	}
	s.mu.Unlock()
	if item.result != nil {
		item.result <- err
	}
}
func (s *Session) write(item outgoing) error {
	if err := item.ctx.Err(); err != nil {
		s.complete(item, err)
		return nil
	}
	deadline := time.Now().Add(s.cfg.WriteTimeout)
	if d, ok := item.ctx.Deadline(); ok && d.Before(deadline) {
		deadline = d
	}
	if err := s.conn.SetWriteDeadline(deadline); err != nil {
		s.complete(item, &OpError{Op: "write", Cause: err})
		return err
	}
	// Only the writer touches Gorilla's deadline field. Cancellation interrupts
	// the underlying socket directly; joining the callback prevents a late expiry
	// from corrupting the following message's deadline.
	canceled := make(chan struct{})
	stop := context.AfterFunc(item.ctx, func() { _ = s.conn.UnderlyingConn().SetWriteDeadline(time.Now()); close(canceled) })
	err := s.conn.WriteMessage(int(item.message.Type), item.message.Data)
	if !stop() {
		<-canceled
	}
	if err != nil {
		if item.ctx.Err() != nil {
			err = item.ctx.Err()
		} else if d, ok := item.ctx.Deadline(); ok && !time.Now().Before(d) {
			// The socket timer can win the race with context's timer.
			err = context.DeadlineExceeded
		}
		wrapped := &OpError{Op: "write", Cause: err}
		s.complete(item, wrapped)
		return wrapped
	}
	s.complete(item, nil)
	return nil
}
func (s *Session) writeControl(frame controlFrame) error {
	return s.conn.WriteControl(frame.kind, frame.data, time.Now().Add(s.cfg.WriteTimeout))
}
func (s *Session) writeLoop() {
	var ticks <-chan time.Time
	if !s.cfg.DisableHeartbeat {
		ticker := time.NewTicker(s.cfg.PingInterval)
		defer ticker.Stop()
		ticks = ticker.C
	}
	for {
		select {
		case <-s.stopped:
			return
		case <-s.closingCh:
			s.finishWrites()
			return
		default:
		}
		select {
		case <-s.stopped:
			return
		case <-s.closingCh:
			s.finishWrites()
			return
		case item := <-s.out:
			if err := s.write(item); err != nil {
				s.terminate(err)
				return
			}
		case frame := <-s.controls:
			if err := s.writeControl(frame); err != nil {
				s.terminate(&OpError{Op: "control", Cause: err})
				return
			}
		case <-ticks:
			if err := s.writeControl(controlFrame{kind: websocket.PingMessage}); err != nil {
				s.terminate(&OpError{Op: "ping", Cause: err})
				return
			}
		}
	}
}
func (s *Session) finishWrites() {
	for {
		select {
		case <-s.stopped:
			return
		case item := <-s.out:
			select {
			case <-s.peerClosed:
				s.complete(item, ErrClosed)
				continue
			default:
			}
			if err := s.write(item); err != nil {
				s.terminate(err)
				return
			}
		default:
			goto closeFrame
		}
	}
closeFrame:
	s.mu.Lock()
	code, reason := s.closeCode, s.closeReason
	s.mu.Unlock()
	if err := s.writeControl(controlFrame{websocket.CloseMessage, websocket.FormatCloseMessage(code, reason)}); err != nil {
		s.terminate(&OpError{Op: "close", Cause: err})
		return
	}
	close(s.closeSent)
}