package main

import (
	"context"
	"encoding/json"
	"log"
	"log/slog"
	"net/http"

	"github.com/bluesky-social/indigo/api/bsky"
	"github.com/bluesky-social/indigo/xrpc"
	"github.com/bluesky-social/jetstream/pkg/client"
	"github.com/bluesky-social/jetstream/pkg/models"
	"github.com/cornelk/hashmap"
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
)

type Set[T comparable] map[T]struct{}

// Data structures
type SubscriberData struct {
	DID      string
	Conn     *websocket.Conn
	ListenTo Set[string]
	Reposts  Set[string]
}

type NotificationMessage struct {
	Liked     bool   `json:"liked"`
	ByDid     string `json:"did"`
	RepostURI string `json:"repost_uri"`
}

// Global state
var (
	subscribers = hashmap.New[string, *SubscriberData]()
	listeningTo = hashmap.New[string, *hashmap.Map[string, *SubscriberData]]()

	likeStream       *client.Client
	subscriberStream *client.Client

	upgrader = websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	logger *slog.Logger
)

func getFollowsDids() []string {
	var dids []string
	subscribers.Range(func(s string, sd *SubscriberData) bool {
		for follow, _ := range sd.ListenTo {
			dids = append(dids, follow)
		}
		return true
	})
	return dids
}

func getSubscriberDids() []string {
	dids := make([]string, 0, subscribers.Len())
	subscribers.Range(func(s string, sd *SubscriberData) bool {
		dids = append(dids, s)
		return true
	})
	return dids
}

func main() {
	logger = slog.Default()

	go likeStreamLoop(logger)
	go subscriberStreamLoop(logger)

	r := mux.NewRouter()
	r.HandleFunc("/subscribe/{did}", handleSubscribe).Methods("GET")

	log.Println("Server starting on :8080")
	if err := http.ListenAndServe(":8080", r); err != nil {
		log.Fatalf("error while serving: %s", err)
	}
}

func handleSubscribe(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	did := vars["did"]

	logger = logger.With("did", did)

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		logger.Error("WebSocket upgrade failed", "error", err)
		return
	}
	defer conn.Close()

	logger.Info("new subscriber")

	pdsURI, err := findUserPDS(r.Context(), did)
	if err != nil {
		logger.Error("cant resolve user pds", "error", err)
		return
	}
	logger = logger.With("pds", pdsURI)

	xrpcClient := &xrpc.Client{
		Host: pdsURI,
	}
	// todo: implement skipping fetching follows and allow specifying users to listen to via websocket
	follows, err := fetchFollows(r.Context(), xrpcClient, did)
	if err != nil {
		logger.Error("error fetching follows", "error", err)
		return
	}
	logger.Info("fetched follows")
	reposts, err := fetchReposts(r.Context(), xrpcClient, did)
	if err != nil {
		logger.Error("error fetching reposts", "error", err)
		return
	}
	logger.Info("fetched reposts")

	sd := &SubscriberData{
		DID:  did,
		Conn: conn,
		// use user follows as default listen to
		ListenTo: follows,
		Reposts:  reposts,
	}

	subscribers.Set(sd.DID, sd)
	for listenDid := range sd.ListenTo {
		listenTo(sd, listenDid)
	}

	updateSubscriberStreamOpts()
	updateLikeStreamOpts()
	// delete subscriber after we are done
	defer func() {
		for listenDid := range sd.ListenTo {
			stopListeningTo(sd.DID, listenDid)
		}
		subscribers.Del(sd.DID)

		updateSubscriberStreamOpts()
		updateLikeStreamOpts()
	}()

	logger.Info("serving subscriber")

	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			logger.Info("WebSocket connection closed", "error", err)
			break
		}
	}
}

func listenTo(sd *SubscriberData, did string) {
	targetDids, _ := listeningTo.GetOrInsert(did, hashmap.New[string, *SubscriberData]())
	targetDids.Insert(sd.DID, sd)
}

func stopListeningTo(subscriberDid, did string) {
	if targetDids, exists := listeningTo.Get(did); exists {
		targetDids.Del(subscriberDid)
	}
}

func getLikeStreamOpts() models.SubscriberOptionsUpdatePayload {
	return models.SubscriberOptionsUpdatePayload{
		WantedCollections: []string{"app.bsky.feed.like"},
		// WantedDIDs:        getFollowsDids(),
	}
}

func getSubscriberStreamOpts() models.SubscriberOptionsUpdatePayload {
	return models.SubscriberOptionsUpdatePayload{
		WantedCollections: []string{"app.bsky.feed.repost", "app.bsky.graph.follow"},
		WantedDIDs:        getSubscriberDids(),
	}
}

func updateLikeStreamOpts() {
	opts := getLikeStreamOpts()
	err := likeStream.SendOptionsUpdate(opts)
	if err != nil {
		logger.Error("couldnt update like stream opts", "error", err)
		return
	}
	logger.Info("updated like stream opts", "requestedDids", len(opts.WantedDIDs))
}

func updateSubscriberStreamOpts() {
	opts := getSubscriberStreamOpts()
	err := subscriberStream.SendOptionsUpdate(opts)
	if err != nil {
		logger.Error("couldnt update subscriber stream opts", "error", err)
		return
	}
	logger.Info("updated subscriber stream opts", "userCount", len(opts.WantedDIDs))
}

func likeStreamLoop(logger *slog.Logger) {
	startJetstreamLoop(logger, &likeStream, "like_tracker", HandleLikeEvent, getLikeStreamOpts)
}

func subscriberStreamLoop(logger *slog.Logger) {
	startJetstreamLoop(logger, &subscriberStream, "subscriber", HandleSubscriberEvent, getSubscriberStreamOpts)
}

func HandleLikeEvent(ctx context.Context, event *models.Event) error {
	if event == nil || event.Commit == nil || len(event.Commit.Record) == 0 {
		return nil
	}

	// skip handling event if its not from a source we are listening to
	targets, exists := listeningTo.Get(event.Did)
	if !exists {
		return nil
	}

	var like bsky.FeedLike
	if err := json.Unmarshal(event.Commit.Record, &like); err != nil {
		logger.Error("failed to unmarshal like", "error", err)
		return nil
	}

	targets.Range(func(s string, sd *SubscriberData) bool {
		for repostURI, _ := range sd.Reposts {
			// (un)liked a post the subscriber reposted
			if like.Subject.Uri == repostURI {
				notification := NotificationMessage{
					Liked:     event.Commit.Operation != models.CommitOperationDelete,
					ByDid:     event.Did,
					RepostURI: repostURI,
				}

				if err := sd.Conn.WriteJSON(notification); err != nil {
					logger.Error("failed to send notification", "subscriber", sd.DID, "error", err)
				}
			}
		}
		return true
	})

	return nil
}

func HandleSubscriberEvent(ctx context.Context, event *models.Event) error {
	if event == nil || event.Commit == nil {
		return nil
	}

	switch event.Commit.Collection {
	case "app.bsky.feed.repost":
		modifySubscribersWithEvent(
			event,
			func(s *SubscriberData, r bsky.FeedRepost) { delete(s.Reposts, r.Subject.Uri) },
			func(s *SubscriberData, r bsky.FeedRepost) {
				s.Reposts[r.Subject.Uri] = struct{}{}
			},
		)
	case "app.bsky.graph.follow":
		modifySubscribersWithEvent(
			event,
			func(s *SubscriberData, r bsky.GraphFollow) {
				delete(s.ListenTo, r.Subject)
				stopListeningTo(s.DID, r.Subject)
			},
			func(s *SubscriberData, r bsky.GraphFollow) {
				s.ListenTo[r.Subject] = struct{}{}
				listenTo(s, r.Subject)
			},
		)
	}

	return nil
}

type ModifyFunc[v any] func(*SubscriberData, v)

func modifySubscribersWithEvent[v any](event *models.Event, onDelete ModifyFunc[v], onUpdate ModifyFunc[v]) error {
	if len(event.Commit.Record) == 0 {
		return nil
	}

	var data v
	if err := json.Unmarshal(event.Commit.Record, &data); err != nil {
		logger.Error("Failed to unmarshal repost", "error", err, "raw", event.Commit.Record)
		return nil
	}

	if subscriber, exists := subscribers.Get(event.Did); exists {
		if event.Commit.Operation == models.CommitOperationDelete {
			onDelete(subscriber, data)
		} else {
			onUpdate(subscriber, data)
		}
	}

	return nil
}