package client

import (
	"log"
	"net"
	"sync"
	"time"

	"github.com/iceBear67/redapricot/client/wire"
)

// WorkerPool manages up to maxConn worker connections and allocates streams
// using the least-loaded strategy (PROTOCOL.md §7.1).
type WorkerPool struct {
	client  *Client
	maxConn int

	mu    sync.Mutex
	conns []*WorkerConn
}

func newWorkerPool(c *Client, maxConn int) *WorkerPool {
	return &WorkerPool{client: c, maxConn: maxConn}
}

// Allocate returns a worker conn and a fresh stream id to place a new stream on.
func (p *WorkerPool) Allocate() (*WorkerConn, int, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	var best *WorkerConn
	bestCount := 0
	for _, wc := range p.conns {
		n := wc.streamCount()
		if best == nil || n < bestCount {
			best = wc
			bestCount = n
		}
	}

	needNew := best == nil || (bestCount > SaturationThreshold && len(p.conns) < p.maxConn)
	if needNew {
		wc, err := p.dialWorker()
		if err != nil {
			if best == nil {
				return nil, 0, err
			}
			log.Printf("worker dial failed, reusing existing conn: %v", err)
		} else {
			p.conns = append(p.conns, wc)
			best = wc
		}
	}
	return best, best.newSid(), nil
}

func (p *WorkerPool) dialWorker() (*WorkerConn, error) {
	fc, err := p.client.dialSession(MagicWorker)
	if err != nil {
		return nil, err
	}
	wc := &WorkerConn{
		pool:    p,
		fc:      fc,
		streams: make(map[int]*Stream),
		nextSid: 1,
	}
	go wc.readLoop()
	log.Printf("opened worker conn (#%d in pool)", len(p.conns)+1)
	return wc, nil
}

func (p *WorkerPool) count() int {
	p.mu.Lock()
	defer p.mu.Unlock()
	return len(p.conns)
}

func (p *WorkerPool) remove(wc *WorkerConn) {
	p.mu.Lock()
	defer p.mu.Unlock()
	for i, c := range p.conns {
		if c == wc {
			p.conns = append(p.conns[:i], p.conns[i+1:]...)
			return
		}
	}
}

func (p *WorkerPool) closeAll() {
	p.mu.Lock()
	conns := append([]*WorkerConn(nil), p.conns...)
	p.mu.Unlock()
	for _, wc := range conns {
		_ = wc.fc.Close()
	}
}

// WorkerConn is one multiplexed worker connection to the hub.
type WorkerConn struct {
	pool *WorkerPool
	fc   *wire.FramedConn

	mu      sync.Mutex
	streams map[int]*Stream
	nextSid int
}

func (wc *WorkerConn) streamCount() int {
	wc.mu.Lock()
	defer wc.mu.Unlock()
	return len(wc.streams)
}

func (wc *WorkerConn) newSid() int {
	wc.mu.Lock()
	defer wc.mu.Unlock()
	sid := wc.nextSid
	wc.nextSid++
	return sid
}

func (wc *WorkerConn) registerStream(sid int, st *Stream) {
	wc.mu.Lock()
	wc.streams[sid] = st
	wc.mu.Unlock()
}

func (wc *WorkerConn) getStream(sid int) *Stream {
	wc.mu.Lock()
	defer wc.mu.Unlock()
	return wc.streams[sid]
}

func (wc *WorkerConn) removeStream(sid int) *Stream {
	wc.mu.Lock()
	defer wc.mu.Unlock()
	st := wc.streams[sid]
	delete(wc.streams, sid)
	return st
}

func (wc *WorkerConn) readLoop() {
	for {
		payload, err := wc.fc.ReadFrame()
		if err != nil {
			break
		}
		r := wire.NewReader(payload)
		ftype, err := r.U8()
		if err != nil {
			continue
		}
		sid, err := r.VarInt()
		if err != nil {
			continue
		}
		switch ftype {
		case MuxData:
			if st := wc.getStream(sid); st != nil {
				st.deliverFromHub(r.Remaining())
			}
		case MuxFin, MuxRst:
			if st := wc.removeStream(sid); st != nil {
				st.shutdown(false)
			}
		default:
			log.Printf("worker: unknown mux type %d", ftype)
		}
	}
	// Connection lost: tear down all streams and drop from pool.
	wc.pool.remove(wc)
	wc.mu.Lock()
	streams := make([]*Stream, 0, len(wc.streams))
	for _, st := range wc.streams {
		streams = append(streams, st)
	}
	wc.streams = make(map[int]*Stream)
	wc.mu.Unlock()
	for _, st := range streams {
		st.shutdown(false)
	}
}

func (wc *WorkerConn) sendSyn(sid int, cid []byte) {
	_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxSyn).VarInt(sid).Bytes(cid).Out())
}

func (wc *WorkerConn) sendData(sid int, data []byte) error {
	return wc.fc.WriteFrame(wire.NewWriter().U8(MuxData).VarInt(sid).Bytes(data).Out())
}

func (wc *WorkerConn) sendFin(sid int) {
	_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxFin).VarInt(sid).Out())
}

func (wc *WorkerConn) sendRst(sid int) {
	_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxRst).VarInt(sid).Out())
}

// Stream bridges one player (via the hub) to one destination connection.
type Stream struct {
	wc      *WorkerConn
	sid     int
	cid     []byte
	mapping Mapping
	srcIP   string
	srcPort int

	mu        sync.Mutex
	dest      net.Conn
	connected bool
	preBuf    []byte
	closed    bool
}

func newStream(wc *WorkerConn, sid int, cid []byte, m Mapping, ip string, port int) *Stream {
	return &Stream{wc: wc, sid: sid, cid: cid, mapping: m, srcIP: ip, srcPort: port}
}

// run dials the destination, optionally writes the PROXY v2 header, flushes any
// buffered hub bytes, then pumps destination -> hub.
func (s *Stream) run() {
	dest, err := net.DialTimeout("tcp", s.mapping.Destination, 10*time.Second)
	if err != nil {
		log.Printf("stream %d: dial %s failed: %v", s.sid, s.mapping.Destination, err)
		s.wc.removeStream(s.sid)
		s.wc.sendRst(s.sid)
		return
	}
	if tcp, ok := dest.(*net.TCPConn); ok {
		_ = tcp.SetNoDelay(true)
	}

	if s.mapping.ProxyProtocol {
		if hdr := s.buildProxyHeader(dest); hdr != nil {
			if _, err := dest.Write(hdr); err != nil {
				log.Printf("stream %d: proxy header write: %v", s.sid, err)
			}
		}
	}

	// Atomically flush pre-connect buffer and enable direct writes.
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		_ = dest.Close()
		return
	}
	s.dest = dest
	if len(s.preBuf) > 0 {
		_, _ = dest.Write(s.preBuf)
		s.preBuf = nil
	}
	s.connected = true
	s.mu.Unlock()

	// destination -> hub
	buf := make([]byte, 32*1024)
	for {
		n, err := dest.Read(buf)
		if n > 0 {
			if werr := s.wc.sendData(s.sid, buf[:n]); werr != nil {
				break
			}
		}
		if err != nil {
			break
		}
	}
	s.shutdown(true)
}

func (s *Stream) buildProxyHeader(dest net.Conn) []byte {
	srcIP := net.ParseIP(s.srcIP)
	if srcIP == nil {
		return nil
	}
	dstTCP, ok := dest.RemoteAddr().(*net.TCPAddr)
	if !ok {
		return nil
	}
	return BuildProxyV2(srcIP, s.srcPort, dstTCP.IP, dstTCP.Port)
}

// deliverFromHub writes bytes coming from the hub to the destination, buffering
// until the destination connection is established.
func (s *Stream) deliverFromHub(data []byte) {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return
	}
	if !s.connected {
		s.preBuf = append(s.preBuf, data...)
		s.mu.Unlock()
		return
	}
	dest := s.dest
	s.mu.Unlock()
	if _, err := dest.Write(data); err != nil {
		s.shutdown(true)
	}
}

// shutdown closes the stream; notifyHub sends a FIN to the hub when true.
func (s *Stream) shutdown(notifyHub bool) {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return
	}
	s.closed = true
	dest := s.dest
	s.mu.Unlock()

	if dest != nil {
		_ = dest.Close()
	}
	s.wc.removeStream(s.sid)
	if notifyHub {
		s.wc.sendFin(s.sid)
	}
}