package main

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"net"
	"sync"
	"time"

	"dragontcp/internal/wire"
)

type streamSession struct {
	sid        wire.SessionID
	target     net.Conn
	targetName string
	maxChunk   int
	maxBuffer  int
	debug      *serverDebug

	mu       sync.Mutex
	notify   chan struct{}
	buf      []byte
	base     uint64
	eof      bool
	closed   bool
	lastSeen time.Time

	upMu       sync.Mutex
	expectedUp uint64
}

func newStreamSession(sid wire.SessionID, target net.Conn, targetName string, maxChunk, maxBuffer int, debug *serverDebug) *streamSession {
	s := &streamSession{
		sid:        sid,
		target:     target,
		targetName: targetName,
		maxChunk:   maxChunk,
		maxBuffer:  maxBuffer,
		debug:      debug,
		notify:     make(chan struct{}),
		lastSeen:   time.Now(),
	}
	go s.readTarget()
	return s
}

func (s *streamSession) signalLocked() {
	close(s.notify)
	s.notify = make(chan struct{})
}

func (s *streamSession) touchLocked() { s.lastSeen = time.Now() }

func (s *streamSession) readTarget() {
	tmp := make([]byte, 64*1024)
	for {
		n, err := s.target.Read(tmp)
		if n > 0 {
			data := append([]byte(nil), tmp[:n]...)
			for len(data) > 0 {
				s.mu.Lock()
				for !s.closed && len(s.buf) >= s.maxBuffer {
					ch := s.notify
					s.mu.Unlock()
					<-ch
					s.mu.Lock()
				}
				if s.closed {
					s.mu.Unlock()
					return
				}
				room := s.maxBuffer - len(s.buf)
				take := len(data)
				if take > room {
					take = room
				}
				s.buf = append(s.buf, data[:take]...)
				data = data[take:]
				s.touchLocked()
				s.signalLocked()
				s.mu.Unlock()
				if s.debug != nil && s.debug.enabled {
					s.debug.bytesDown.Add(uint64(take))
				}
			}
		}
		if err != nil {
			s.mu.Lock()
			if !s.closed {
				s.eof = true
				s.touchLocked()
				s.signalLocked()
			}
			s.mu.Unlock()
			return
		}
	}
}

func (s *streamSession) ackLocked(offset uint64) {
	if offset <= s.base {
		return
	}
	end := s.base + uint64(len(s.buf))
	if offset > end {
		offset = end
	}
	drop := int(offset - s.base)
	if drop <= 0 {
		return
	}
	s.buf = s.buf[drop:]
	s.base = offset
	if len(s.buf) == 0 {
		s.buf = nil
	} else if cap(s.buf) > 4*len(s.buf) && cap(s.buf) > 1024*1024 {
		compact := append([]byte(nil), s.buf...)
		s.buf = compact
	}
	s.signalLocked()
}

func (s *streamSession) ack(offset uint64) {
	s.mu.Lock()
	s.ackLocked(offset)
	s.touchLocked()
	s.mu.Unlock()
}

func (s *streamSession) readAt(offset uint64, limit int, wait time.Duration) ([]byte, byte, error) {
	if limit < 1 || limit > s.maxChunk {
		return nil, wire.StatusError, fmt.Errorf("invalid download limit %d", limit)
	}
	deadline := time.Now().Add(wait)
	firstDataAt := time.Time{}

	for {
		s.mu.Lock()
		s.touchLocked()
		if offset < s.base {
			s.mu.Unlock()
			return nil, wire.StatusError, fmt.Errorf("download offset %d was already acknowledged (base=%d)", offset, s.base)
		}
		rel64 := offset - s.base
		if rel64 <= uint64(len(s.buf)) {
			rel := int(rel64)
			available := len(s.buf) - rel
			if available > 0 {
				if firstDataAt.IsZero() {
					firstDataAt = time.Now()
				}
				// Coalesce tiny target reads briefly. This prevents a 1-2 byte
				// producer read from becoming a permanent tiny tunnel record.
				if available < limit && !s.eof && wait > 0 && time.Since(firstDataAt) < 2*time.Millisecond {
					ch := s.notify
					s.mu.Unlock()
					select {
					case <-ch:
					case <-time.After(2 * time.Millisecond):
					}
					continue
				}
				n := available
				if n > limit {
					n = limit
				}
				out := append([]byte(nil), s.buf[rel:rel+n]...)
				s.mu.Unlock()
				return out, wire.StatusData, nil
			}
			if s.eof || s.closed {
				s.mu.Unlock()
				return nil, wire.StatusEOF, nil
			}
		} else {
			s.mu.Unlock()
			return nil, wire.StatusError, fmt.Errorf("download offset %d is beyond buffered stream end %d", offset, s.base+uint64(len(s.buf)))
		}

		if wait <= 0 || time.Now().After(deadline) {
			s.mu.Unlock()
			return nil, wire.StatusWait, nil
		}
		ch := s.notify
		remaining := time.Until(deadline)
		s.mu.Unlock()
		select {
		case <-ch:
		case <-time.After(remaining):
			return nil, wire.StatusWait, nil
		}
	}
}

func (s *streamSession) upload(offset uint64, data []byte) error {
	if len(data) == 0 || len(data) > s.maxChunk {
		return fmt.Errorf("invalid upload size %d", len(data))
	}
	s.upMu.Lock()
	defer s.upMu.Unlock()

	if offset < s.expectedUp {
		// Idempotent retry after a lost ACK.
		if offset+uint64(len(data)) <= s.expectedUp {
			return nil
		}
		return fmt.Errorf("overlapping upload retry at %d", offset)
	}
	if offset != s.expectedUp {
		return fmt.Errorf("upload gap: got %d expected %d", offset, s.expectedUp)
	}
	if _, err := s.target.Write(data); err != nil {
		return err
	}
	s.expectedUp += uint64(len(data))
	s.mu.Lock()
	s.touchLocked()
	s.mu.Unlock()
	if s.debug != nil && s.debug.enabled {
		s.debug.bytesUp.Add(uint64(len(data)))
		s.debug.pushRecords.Add(1)
	}
	return nil
}

func (s *streamSession) close() {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return
	}
	s.closed = true
	s.signalLocked()
	s.mu.Unlock()
	_ = s.target.Close()
}

type streamManager struct {
	mu       sync.RWMutex
	sessions map[string]*streamSession
	timeout  time.Duration
	debug    *serverDebug
}

func sidKey(sid wire.SessionID) string { return string(sid[:]) }

func newStreamManager(timeout time.Duration, debug *serverDebug) *streamManager {
	m := &streamManager{sessions: make(map[string]*streamSession), timeout: timeout, debug: debug}
	go m.cleanupLoop()
	return m
}

func (m *streamManager) get(sid wire.SessionID) *streamSession {
	m.mu.RLock()
	s := m.sessions[sidKey(sid)]
	m.mu.RUnlock()
	return s
}

func (m *streamManager) addOrGet(sid wire.SessionID, s *streamSession) (*streamSession, bool) {
	key := sidKey(sid)
	m.mu.Lock()
	if old := m.sessions[key]; old != nil {
		m.mu.Unlock()
		s.close()
		return old, false
	}
	m.sessions[key] = s
	m.mu.Unlock()
	return s, true
}

func (m *streamManager) remove(sid wire.SessionID) {
	key := sidKey(sid)
	m.mu.Lock()
	s := m.sessions[key]
	delete(m.sessions, key)
	m.mu.Unlock()
	if s != nil {
		s.close()
		if m.debug != nil && m.debug.enabled {
			m.debug.sessionsClosed.Add(1)
			m.debug.activeSessions.Add(-1)
		}
	}
}

func (m *streamManager) count() int { m.mu.RLock(); n := len(m.sessions); m.mu.RUnlock(); return n }

func (m *streamManager) cleanupLoop() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	for range ticker.C {
		cutoff := time.Now().Add(-m.timeout)
		var stale []wire.SessionID
		m.mu.RLock()
		for _, s := range m.sessions {
			s.mu.Lock()
			last := s.lastSeen
			closed := s.closed
			sid := s.sid
			s.mu.Unlock()
			if closed || last.Before(cutoff) {
				stale = append(stale, sid)
			}
		}
		m.mu.RUnlock()
		for _, sid := range stale {
			m.remove(sid)
		}
	}
}

func parseProbe(payload []byte) (kind byte, value int, token string, err error) {
	if len(payload) < 11 || !bytes.Equal(payload[:4], wire.ProbeMagic[:]) {
		return 0, 0, "", fmt.Errorf("bad probe payload")
	}
	kind = payload[4]
	tl := int(binary.BigEndian.Uint16(payload[5:7]))
	value = int(binary.BigEndian.Uint32(payload[7:11]))
	if 11+tl > len(payload) {
		return 0, 0, "", fmt.Errorf("bad probe token length")
	}
	token = string(payload[11 : 11+tl])
	return
}

func probePattern(n int) []byte {
	out := make([]byte, n)
	for i := range out {
		out[i] = byte((i*31 + 17) & 0xff)
	}
	return out
}

func parseOpen(payload []byte) (token, host string, port int, err error) {
	if len(payload) < 6 {
		return "", "", 0, fmt.Errorf("bad OPEN payload")
	}
	tl := int(binary.BigEndian.Uint16(payload[0:2]))
	hl := int(binary.BigEndian.Uint16(payload[2:4]))
	port = int(binary.BigEndian.Uint16(payload[4:6]))
	if port < 1 || 6+tl+hl != len(payload) {
		return "", "", 0, fmt.Errorf("bad OPEN lengths")
	}
	token = string(payload[6 : 6+tl])
	host = string(payload[6+tl:])
	if host == "" {
		return "", "", 0, fmt.Errorf("empty target host")
	}
	return
}

func processWireRequest(conn net.Conn, req wire.Request, token string, allowPrivate bool, cache *dnsCache, tcpBuffer int, manager *streamManager, maxChunk, maxBuffer int, pollWait time.Duration, debug *serverDebug) error {
	switch req.Mode {
	case wire.ModeProbe:
		kind, value, supplied, err := parseProbe(req.Payload)
		if err != nil {
			return wire.WriteResponse(conn, wire.StatusError, []byte(err.Error()))
		}
		if !tokenEqual(supplied, token) {
			return wire.WriteResponse(conn, wire.StatusError, []byte("authentication failed"))
		}
		switch kind {
		case wire.ProbeUpload:
			if len(req.Payload) > maxChunk {
				return wire.WriteResponse(conn, wire.StatusError, []byte("probe too large"))
			}
			return wire.WriteResponse(conn, wire.StatusOK, nil)
		case wire.ProbeDownload:
			if value < 1 || value > maxChunk {
				return wire.WriteResponse(conn, wire.StatusError, []byte("probe too large"))
			}
			return wire.WriteMaskedResponse(conn, wire.StatusData, probePattern(value), req.Session, wire.ModeProbe, req.Seq)
		case wire.ProbeKeepalive:
			return wire.WriteResponse(conn, wire.StatusOK, nil)
		case wire.ProbeBatch:
			count := value
			if count < 1 {
				count = 1
			}
			if count > 16 {
				count = 16
			}
			for i := 0; i < count; i++ {
				data := probePattern(32)
				if err := wire.WriteMaskedResponse(conn, wire.StatusData, data, req.Session, wire.ModeProbe, req.Seq+uint64(i)); err != nil {
					return err
				}
			}
			return nil
		default:
			return wire.WriteResponse(conn, wire.StatusError, []byte("unknown probe kind"))
		}

	case wire.ModeOpen:
		supplied, host, port, err := parseOpen(req.Payload)
		if err != nil {
			return wire.WriteResponse(conn, wire.StatusError, []byte(err.Error()))
		}
		if !tokenEqual(supplied, token) {
			return wire.WriteResponse(conn, wire.StatusError, []byte("authentication failed"))
		}
		if old := manager.get(req.Session); old != nil {
			body := make([]byte, 4)
			binary.BigEndian.PutUint32(body, uint32(maxChunk))
			return wire.WriteMaskedResponse(conn, wire.StatusOK, body, req.Session, wire.ModeOpen, req.Seq)
		}
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		target, err := dialTarget(ctx, host, port, allowPrivate, cache, tcpBuffer)
		cancel()
		if err != nil {
			return wire.WriteResponse(conn, wire.StatusError, []byte(err.Error()))
		}
		session := newStreamSession(req.Session, target, fmt.Sprintf("%s:%d", host, port), maxChunk, maxBuffer, debug)
		_, created := manager.addOrGet(req.Session, session)
		if created && debug != nil && debug.enabled {
			debug.sessionsOpened.Add(1)
			debug.activeSessions.Add(1)
			debug.logf("SESSION OPEN sid=%x target=%s:%d active_sessions=%d", req.Session[:4], host, port, manager.count())
		}
		body := make([]byte, 4)
		binary.BigEndian.PutUint32(body, uint32(maxChunk))
		return wire.WriteMaskedResponse(conn, wire.StatusOK, body, req.Session, wire.ModeOpen, req.Seq)

	case wire.ModeUpload:
		s := manager.get(req.Session)
		if s == nil {
			return wire.WriteResponse(conn, wire.StatusError, []byte("unknown session"))
		}
		if len(req.Payload) > maxChunk {
			return wire.WriteResponse(conn, wire.StatusError, []byte("upload too large"))
		}
		if err := s.upload(req.Seq, req.Payload); err != nil {
			return wire.WriteResponse(conn, wire.StatusError, []byte(err.Error()))
		}
		return wire.WriteResponse(conn, wire.StatusOK, nil)

	case wire.ModeDownload:
		s := manager.get(req.Session)
		if s == nil {
			return wire.WriteResponse(conn, wire.StatusError, []byte("unknown session"))
		}
		if len(req.Payload) != 14 {
			return wire.WriteResponse(conn, wire.StatusError, []byte("bad download request"))
		}
		ack := binary.BigEndian.Uint64(req.Payload[0:8])
		limit := int(binary.BigEndian.Uint32(req.Payload[8:12]))
		count := int(binary.BigEndian.Uint16(req.Payload[12:14]))
		if limit < 1 {
			limit = 1
		}
		if limit > maxChunk {
			limit = maxChunk
		}
		if count < 1 {
			count = 1
		}
		if count > 256 {
			count = 256
		}
		s.ack(ack)
		offset := req.Seq
		if debug != nil && debug.enabled {
			debug.pullRequests.Add(1)
		}
		for i := 0; i < count; i++ {
			wait := time.Duration(0)
			if i == 0 {
				wait = pollWait
			}
			data, status, err := s.readAt(offset, limit, wait)
			if err != nil {
				return wire.WriteResponse(conn, wire.StatusError, []byte(err.Error()))
			}
			switch status {
			case wire.StatusData:
				if debug != nil && debug.enabled {
					debug.dataRecords.Add(1)
				}
				if err := wire.WriteMaskedResponse(conn, wire.StatusData, data, req.Session, wire.ModeDownload, offset); err != nil {
					return err
				}
				offset += uint64(len(data))
			case wire.StatusWait:
				if debug != nil && debug.enabled {
					debug.waitRecords.Add(1)
				}
				return wire.WriteResponse(conn, wire.StatusWait, nil)
			case wire.StatusEOF:
				return wire.WriteResponse(conn, wire.StatusEOF, nil)
			default:
				return wire.WriteResponse(conn, wire.StatusError, []byte("invalid session read status"))
			}
		}
		return nil

	case wire.ModeClose:
		manager.remove(req.Session)
		return wire.WriteResponse(conn, wire.StatusOK, nil)
	default:
		return wire.WriteResponse(conn, wire.StatusError, []byte("unknown mode"))
	}
}