package main

import (
	"encoding/json"
	"log"
	"net"
	"net/http"
	"regexp"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

const (
	maxFrameBytes  = 64 << 10 // SDP is 2-8 KiB; encrypted blobs stay well under this
	writeWait      = 10 * time.Second
	pongWait       = 60 * time.Second
	pingPeriod     = 25 * time.Second
	sendQueueDepth = 64
)

var (
	roomRe = regexp.MustCompile(`^[A-Za-z0-9_-]{43}$`) // base64url of 32 bytes
	peerRe = regexp.MustCompile(`^[a-f0-9]{32}$`)      // hex of 16 bytes
)

type Hub struct {
	cfg      *Config
	upgrader websocket.Upgrader

	mu    sync.Mutex
	rooms map[string]map[string]*client
	perIP map[string]int
}

type client struct {
	hub    *Hub
	ws     *websocket.Conn
	send   chan []byte
	ip     string
	bucket *bucket

	mu   sync.Mutex
	room string
	peer string
}

func newHub(cfg *Config) *Hub {
	h := &Hub{
		cfg:   cfg,
		rooms: map[string]map[string]*client{},
		perIP: map[string]int{},
	}
	h.upgrader = websocket.Upgrader{
		ReadBufferSize:  4096,
		WriteBufferSize: 4096,
		CheckOrigin: func(r *http.Request) bool {
			return cfg.originAllowed(r.Header.Get("Origin"))
		},
	}
	return h
}

func (h *Hub) stats() (rooms, peers int) {
	h.mu.Lock()
	defer h.mu.Unlock()
	for _, m := range h.rooms {
		peers += len(m)
	}
	return len(h.rooms), peers
}

func (h *Hub) handleWS(w http.ResponseWriter, r *http.Request) {
	ip, _, _ := net.SplitHostPort(r.RemoteAddr)

	h.mu.Lock()
	if h.perIP[ip] >= h.cfg.MaxConnsPerIP {
		h.mu.Unlock()
		http.Error(w, "too many connections", http.StatusTooManyRequests)
		return
	}
	h.perIP[ip]++
	h.mu.Unlock()

	ws, err := h.upgrader.Upgrade(w, r, nil)
	if err != nil {
		h.dropIP(ip)
		return
	}
	c := &client{
		hub:    h,
		ws:     ws,
		send:   make(chan []byte, sendQueueDepth),
		ip:     ip,
		bucket: newBucket(30, 60),
	}
	go c.writePump()
	c.readPump()
}

func (h *Hub) dropIP(ip string) {
	h.mu.Lock()
	if h.perIP[ip]--; h.perIP[ip] <= 0 {
		delete(h.perIP, ip)
	}
	h.mu.Unlock()
}

// wire types

type inMsg struct {
	T       string          `json:"t"`
	Room    string          `json:"room,omitempty"`
	Peer    string          `json:"peer,omitempty"`
	To      string          `json:"to,omitempty"`
	Payload json.RawMessage `json:"payload,omitempty"`
}

type outMsg struct {
	T       string          `json:"t"`
	Peer    string          `json:"peer,omitempty"`
	Peers   []string        `json:"peers,omitzero"`
	From    string          `json:"from,omitempty"`
	Payload json.RawMessage `json:"payload,omitempty"`
	Code    string          `json:"code,omitempty"`
}

func (c *client) sendJSON(m outMsg) {
	b, _ := json.Marshal(m)
	select {
	case c.send <- b:
	default:
		// Slow consumer: drop the connection rather than block the room.
		c.ws.Close()
	}
}

func (c *client) readPump() {
	defer c.teardown()
	c.ws.SetReadLimit(maxFrameBytes)
	c.ws.SetReadDeadline(time.Now().Add(pongWait))
	c.ws.SetPongHandler(func(string) error {
		c.ws.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})
	for {
		_, data, err := c.ws.ReadMessage()
		if err != nil {
			return
		}
		if !c.bucket.take() {
			c.sendJSON(outMsg{T: "error", Code: "rate-limited"})
			continue
		}
		var m inMsg
		if err := json.Unmarshal(data, &m); err != nil {
			c.sendJSON(outMsg{T: "error", Code: "bad-message"})
			continue
		}
		switch m.T {
		case "join":
			c.handleJoin(m)
		case "signal":
			c.handleSignal(m)
		case "leave":
			return // teardown broadcasts peer-left
		case "ping":
			c.sendJSON(outMsg{T: "pong"})
		default:
			c.sendJSON(outMsg{T: "error", Code: "bad-message"})
		}
	}
}

func (c *client) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.ws.Close()
	}()
	for {
		select {
		case b, ok := <-c.send:
			c.ws.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				c.ws.WriteMessage(websocket.CloseMessage, nil)
				return
			}
			if err := c.ws.WriteMessage(websocket.TextMessage, b); err != nil {
				return
			}
		case <-ticker.C:
			c.ws.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.ws.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

func (c *client) handleJoin(m inMsg) {
	if !roomRe.MatchString(m.Room) || !peerRe.MatchString(m.Peer) {
		c.sendJSON(outMsg{T: "error", Code: "bad-message"})
		return
	}
	c.mu.Lock()
	alreadyJoined := c.room != ""
	c.mu.Unlock()
	if alreadyJoined {
		c.sendJSON(outMsg{T: "error", Code: "bad-message"})
		return
	}

	h := c.hub
	h.mu.Lock()
	room := h.rooms[m.Room]
	if room == nil {
		if len(h.rooms) >= h.cfg.MaxRooms {
			h.mu.Unlock()
			c.sendJSON(outMsg{T: "error", Code: "too-many-rooms"})
			return
		}
		room = map[string]*client{}
		h.rooms[m.Room] = room
	}
	// Same peerId rejoining (e.g. zombie socket after mobile sleep): replace.
	if old := room[m.Peer]; old != nil {
		delete(room, m.Peer)
		defer old.ws.Close()
	} else if len(room) >= h.cfg.MaxRoomPeers {
		h.mu.Unlock()
		c.sendJSON(outMsg{T: "error", Code: "room-full"})
		return
	}
	peers := make([]string, 0, len(room))
	others := make([]*client, 0, len(room))
	for id, cl := range room {
		peers = append(peers, id)
		others = append(others, cl)
	}
	room[m.Peer] = c
	h.mu.Unlock()

	c.mu.Lock()
	c.room, c.peer = m.Room, m.Peer
	c.mu.Unlock()

	c.sendJSON(outMsg{T: "peers", Peers: peers})
	for _, o := range others {
		o.sendJSON(outMsg{T: "peer-joined", Peer: m.Peer})
	}
	log.Printf("join: room peers=%d (total rooms=%d)", len(peers)+1, roomCount(h))
}

func (c *client) handleSignal(m inMsg) {
	c.mu.Lock()
	room, me := c.room, c.peer
	c.mu.Unlock()
	if room == "" {
		c.sendJSON(outMsg{T: "error", Code: "not-joined"})
		return
	}
	if !peerRe.MatchString(m.To) || len(m.Payload) == 0 {
		c.sendJSON(outMsg{T: "error", Code: "bad-message"})
		return
	}
	h := c.hub
	h.mu.Lock()
	target := h.rooms[room][m.To]
	h.mu.Unlock()
	if target == nil {
		return // peer left; silently drop
	}
	// `from` is stamped from the joined connection — clients cannot spoof it.
	target.sendJSON(outMsg{T: "signal", From: me, Payload: m.Payload})
}

func (c *client) teardown() {
	c.ws.Close()
	c.mu.Lock()
	room, peer := c.room, c.peer
	c.room, c.peer = "", ""
	c.mu.Unlock()

	h := c.hub
	h.dropIP(c.ip)
	if room == "" {
		return
	}
	h.mu.Lock()
	var others []*client
	if m := h.rooms[room]; m != nil && m[peer] == c {
		delete(m, peer)
		for _, cl := range m {
			others = append(others, cl)
		}
		if len(m) == 0 {
			delete(h.rooms, room)
		}
	}
	h.mu.Unlock()
	for _, o := range others {
		o.sendJSON(outMsg{T: "peer-left", Peer: peer})
	}
}

func roomCount(h *Hub) int {
	h.mu.Lock()
	defer h.mu.Unlock()
	return len(h.rooms)
}