// bookd serves Injective order books over REST and WebSocket. One poller runs
// per subscribed market and fans out snapshots and level diffs to every
// WebSocket client, so N clients cost one upstream request per tick.
//
//	go run ./cmd/bookd -addr :8080
//	curl localhost:8080/v1/book/<market_id>?kind=spot
//	websocat ws://localhost:8080/v1/stream/<market_id>?kind=spot
package main

import (
	"context"
	"encoding/json"
	"flag"
	"log"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/coder/websocket"
	"github.com/coder/websocket/wsjson"

	"github.com/edwardtay/injective-booklab/book"
	"github.com/edwardtay/injective-booklab/client"
	"github.com/edwardtay/injective-booklab/dec"
)

type msg struct {
	Type    string        `json:"type"` // "snapshot" or "diff"
	Seq     uint64        `json:"seq"`
	Book    *book.Book    `json:"book,omitempty"`
	Changes []book.Change `json:"changes,omitempty"`
}

type feed struct {
	mu   sync.Mutex
	subs map[chan msg]struct{}
	last *book.Book
	stop context.CancelFunc
}

type hub struct {
	c     *client.Client
	every time.Duration
	mu    sync.Mutex
	feeds map[string]*feed
}

func (h *hub) subscribe(kind, id string) (chan msg, *book.Book, func()) {
	h.mu.Lock()
	f, ok := h.feeds[id]
	if !ok {
		ctx, cancel := context.WithCancel(context.Background())
		f = &feed{subs: map[chan msg]struct{}{}, stop: cancel}
		h.feeds[id] = f
		go h.poll(ctx, f, kind, id)
	}
	h.mu.Unlock()
	ch := make(chan msg, 16)
	f.mu.Lock()
	f.subs[ch] = struct{}{}
	last := f.last
	f.mu.Unlock()
	return ch, last, func() {
		f.mu.Lock()
		delete(f.subs, ch)
		empty := len(f.subs) == 0
		f.mu.Unlock()
		if empty {
			h.mu.Lock()
			if h.feeds[id] == f {
				delete(h.feeds, id)
				f.stop()
			}
			h.mu.Unlock()
		}
	}
}

func (h *hub) poll(ctx context.Context, f *feed, kind, id string) {
	t := time.NewTicker(h.every)
	defer t.Stop()
	for {
		rctx, cancel := context.WithTimeout(ctx, 5*time.Second)
		snap, err := h.c.Orderbook(rctx, kind, id, 50)
		cancel()
		if err == nil {
			f.mu.Lock()
			var m msg
			switch {
			case f.last == nil:
				m = msg{Type: "snapshot", Seq: snap.Book.Seq, Book: snap.Book}
			case snap.Book.Seq != f.last.Seq:
				m = msg{Type: "diff", Seq: snap.Book.Seq, Changes: book.Diff(f.last, snap.Book)}
			}
			f.last = snap.Book
			for ch := range f.subs {
				if m.Type == "" {
					break // unchanged seq: nothing to send
				}
				select {
				case ch <- m:
				default: // slow consumer: drop; it resyncs from the next snapshot request
				}
			}
			f.mu.Unlock()
		} else if ctx.Err() == nil {
			log.Printf("poll %s: %v", id, err)
		}
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		}
	}
}

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(v)
}

func main() {
	addr := flag.String("addr", ":8080", "listen address")
	every := flag.Duration("every", time.Second, "poll interval per market")
	flag.Parse()
	h := &hub{c: client.New(), every: *every, feeds: map[string]*feed{}}
	log.Printf("bookd listening on %s", *addr)
	srv := &http.Server{Addr: *addr, Handler: routes(h), ReadHeaderTimeout: 5 * time.Second}
	log.Fatal(srv.ListenAndServe())
}

func routes(h *hub) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /v1/markets", func(w http.ResponseWriter, r *http.Request) {
		ms, err := h.c.Markets(r.Context())
		if err != nil {
			writeJSON(w, http.StatusBadGateway, map[string]string{"error": err.Error()})
			return
		}
		writeJSON(w, http.StatusOK, ms)
	})
	mux.HandleFunc("GET /v1/book/{id}", func(w http.ResponseWriter, r *http.Request) {
		kind := r.URL.Query().Get("kind")
		if kind != "derivative" {
			kind = "spot"
		}
		limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
		if limit <= 0 || limit > 500 {
			limit = 50
		}
		s, err := h.c.Orderbook(r.Context(), kind, r.PathValue("id"), limit)
		if err != nil {
			writeJSON(w, http.StatusBadGateway, map[string]string{"error": err.Error()})
			return
		}
		writeJSON(w, http.StatusOK, s)
	})
	mux.HandleFunc("GET /v1/sim/{id}", func(w http.ResponseWriter, r *http.Request) {
		q := r.URL.Query()
		side, err := book.ParseSide(q.Get("side"))
		qty, err2 := dec.Parse(q.Get("qty"))
		if err != nil || err2 != nil {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "side=buy|sell and qty=<decimal> are required"})
			return
		}
		kind := q.Get("kind")
		if kind != "derivative" {
			kind = "spot"
		}
		s, err := h.c.Orderbook(r.Context(), kind, r.PathValue("id"), 500)
		if err != nil {
			writeJSON(w, http.StatusBadGateway, map[string]string{"error": err.Error()})
			return
		}
		fee, _ := dec.Parse(q.Get("fee"))
		writeJSON(w, http.StatusOK, book.SimulateTaker(s.Book, side, qty, dec.Zero(), fee))
	})
	mux.HandleFunc("GET /v1/health", func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, http.StatusOK, h.c.Stats())
	})
	mux.HandleFunc("GET /v1/stream/{id}", func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{InsecureSkipVerify: true})
		if err != nil {
			return
		}
		defer conn.CloseNow()
		kind := r.URL.Query().Get("kind")
		if kind != "derivative" {
			kind = "spot"
		}
		ch, last, unsub := h.subscribe(kind, r.PathValue("id"))
		defer unsub()
		ctx := conn.CloseRead(r.Context())
		if last != nil {
			if wsjson.Write(ctx, conn, msg{Type: "snapshot", Seq: last.Seq, Book: last}) != nil {
				return
			}
		}
		for {
			select {
			case <-ctx.Done():
				return
			case m := <-ch:
				wctx, cancel := context.WithTimeout(ctx, 5*time.Second)
				err := wsjson.Write(wctx, conn, m)
				cancel()
				if err != nil {
					return
				}
			}
		}
	})
	return mux
}