package ws

import (
	"errors"
	"io"
	"net"
	"net/url"
	"sync"
	"time"
)

// StreamConn adapts a WebSocket to net.Conn: Read returns the payloads of
// binary or text messages in order, Write sends one binary message. A P2P
// stream is a byte stream, so message boundaries are irrelevant above this.
type StreamConn struct {
	c      *Conn
	local  net.Addr
	remote net.Addr

	mu  sync.Mutex
	buf []byte
	err error
}

// DialStream opens a WebSocket and returns it as a net.Conn.
func DialStream(rawURL string) (*StreamConn, error) {
	c, err := Dial(rawURL)
	if err != nil {
		return nil, err
	}
	host := "ws"
	if u, err := url.Parse(rawURL); err == nil && u.Host != "" {
		host = u.Host
	}
	return &StreamConn{
		c:      c,
		local:  streamAddr("stream:local"),
		remote: streamAddr(host),
	}, nil
}

// Read implements net.Conn.
func (s *StreamConn) Read(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}
	for {
		s.mu.Lock()
		if len(s.buf) > 0 {
			n := copy(p, s.buf)
			s.buf = s.buf[n:]
			s.mu.Unlock()
			return n, nil
		}
		err := s.err
		s.mu.Unlock()
		if err != nil {
			return 0, err
		}
		msg, err := s.c.ReadMessage()
		if err != nil {
			if errors.Is(err, io.EOF) {
				err = io.EOF
			}
			s.mu.Lock()
			s.err = err
			s.mu.Unlock()
			return 0, err
		}
		if len(msg) == 0 {
			continue
		}
		s.mu.Lock()
		s.buf = msg
		s.mu.Unlock()
	}
}

// Write implements net.Conn.
func (s *StreamConn) Write(p []byte) (int, error) {
	if err := s.c.WriteBinary(p); err != nil {
		return 0, err
	}
	return len(p), nil
}

// Close implements net.Conn.
func (s *StreamConn) Close() error { return s.c.Close() }

// LocalAddr implements net.Conn.
func (s *StreamConn) LocalAddr() net.Addr { return s.local }

// RemoteAddr implements net.Conn.
func (s *StreamConn) RemoteAddr() net.Addr { return s.remote }

// SetDeadline implements net.Conn.
func (s *StreamConn) SetDeadline(t time.Time) error { return s.c.SetDeadline(t) }

// SetReadDeadline implements net.Conn.
func (s *StreamConn) SetReadDeadline(t time.Time) error { return s.c.SetReadDeadline(t) }

// SetWriteDeadline implements net.Conn.
func (s *StreamConn) SetWriteDeadline(t time.Time) error { return s.c.SetWriteDeadline(t) }

type streamAddr string

func (a streamAddr) Network() string { return "ws" }
func (a streamAddr) String() string  { return string(a) }