package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"strings"
	"time"

	"nhooyr.io/websocket"
)

// Ingestor connects to the Bybit V5 WebSocket, receives raw JSON frames,
// extracts the topic prefix, and routes messages to registered MessageHandlers.
type Ingestor struct {
	cfg      Config
	handlers map[string]MessageHandler // topic prefix -> handler
}

// NewIngestor creates an Ingestor with registered stream handlers.
func NewIngestor(cfg Config, handlers map[string]MessageHandler) *Ingestor {
	return &Ingestor{
		cfg:      cfg,
		handlers: handlers,
	}
}

// Run connects to the WebSocket and routes messages until ctx is cancelled.
func (ing *Ingestor) Run(ctx context.Context) {
	log.Println("[ingestor] Starting multi-stream WebSocket ingestor...")
	defer log.Println("[ingestor] Ingestor stopped.")

	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		if err := ing.connectAndConsume(ctx); err != nil {
			if ctx.Err() != nil {
				return // Context cancelled
			}
			log.Printf("[ingestor] Connection error: %v. Reconnecting in 5s...", err)
			select {
			case <-time.After(5 * time.Second):
			case <-ctx.Done():
				return
			}
		}
	}
}

func (ing *Ingestor) connectAndConsume(ctx context.Context) error {
	url := ing.cfg.WebSocketURL
	log.Printf("[ingestor] Connecting to %s ...", url)

	conn, _, err := websocket.Dial(ctx, url, nil)
	if err != nil {
		return fmt.Errorf("dial: %w", err)
	}
	defer conn.CloseNow()

	// 2 MB read limit for large orderbook snapshots or trade batches
	conn.SetReadLimit(2 << 20)

	// Collect all topics from registered handlers
	var allTopics []string
	for _, h := range ing.handlers {
		allTopics = append(allTopics, h.Topics()...)
	}

	if len(allTopics) == 0 {
		return fmt.Errorf("no topics registered to subscribe")
	}

	subMsg := map[string]interface{}{
		"op":   "subscribe",
		"args": allTopics,
	}
	subJSON, _ := json.Marshal(subMsg)

	if err := conn.Write(ctx, websocket.MessageText, subJSON); err != nil {
		return fmt.Errorf("subscribe: %w", err)
	}
	log.Printf("[ingestor] Subscribed to topics: %v", allTopics)

	// Ping ticker to keep Bybit connection alive
	pingTicker := time.NewTicker(20 * time.Second)
	defer pingTicker.Stop()

	go func() {
		for {
			select {
			case <-pingTicker.C:
				pingMsg := []byte(`{"op":"ping"}`)
				if err := conn.Write(ctx, websocket.MessageText, pingMsg); err != nil {
					return
				}
			case <-ctx.Done():
				return
			}
		}
	}()

	// Read loop
	for {
		select {
		case <-ctx.Done():
			conn.Close(websocket.StatusNormalClosure, "shutting down")
			return nil
		default:
		}

		_, data, err := conn.Read(ctx)
		if err != nil {
			return fmt.Errorf("read: %w", err)
		}

		ing.handleMessage(data)
	}
}

type topicHeader struct {
	Topic string `json:"topic"`
}

func (ing *Ingestor) handleMessage(data []byte) {
	var header topicHeader
	if err := json.Unmarshal(data, &header); err != nil || header.Topic == "" {
		return // pong, sub ack, or malformed message
	}

	prefix := extractTopicPrefix(header.Topic)
	if handler, ok := ing.handlers[prefix]; ok {
		handler.HandleMessage(data)
	}
}

// extractTopicPrefix extracts the topic family name (e.g. "publicTrade", "tickers", "kline", "orderbook", "allLiquidation").
func extractTopicPrefix(topic string) string {
	parts := strings.Split(topic, ".")
	if len(parts) > 0 {
		return parts[0]
	}
	return topic
}