package ws

import (
	"context"
	"encoding/json"
	"errors"
	"homestead/homestead_gateway/util/config"
	"log/slog"
	"net"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
)

func NewWsGateway(cfg config.GatewayConfig, logger *slog.Logger, closefn func() error, modH ModHandler, botH BotHandler) *WebsocketGateway {
	return &WebsocketGateway{
		logger:  logger,
		closefn: closefn,
		apiKey:  cfg.Websocket,
		upgrader: websocket.Upgrader{
			ReadBufferSize:  1024,
			WriteBufferSize: 1024,
			CheckOrigin: func(r *http.Request) bool {
				return true // local by default; change for production
			},
		},
		conns:         make(map[*websocket.Conn]connMetadata),
		bodySizeBytes: int64(cfg.BodySize) * 1024 * 1024,
		port:          cfg.HttpPort,
		modHandler:    modH,
		botHandler:    botH,
	}
}

func (g *WebsocketGateway) Serve(ctx context.Context, listenAddr string) error {
	mux := http.NewServeMux()
	mux.HandleFunc("/ws", g.handleWS)
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		_, _ = w.Write([]byte("ok"))
	})

	srv := &http.Server{
		Addr:        listenAddr,
		Handler:     loggingMiddleware(g.logger, mux),
		BaseContext: func(l net.Listener) context.Context { return ctx },
	}

	errCh := make(chan error, 1)
	go func() {
		g.logger.Info("ws gateway listening", "addr", listenAddr)
		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			errCh <- err
		}
		close(errCh)
	}()

	select {
	case <-ctx.Done():
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		g.logger.Info("shutting down http server")
		_ = srv.Shutdown(shutdownCtx)
		g.closeAll()
		return nil
	case err := <-errCh:
		return err
	}
}

func (g *WebsocketGateway) handleWS(w http.ResponseWriter, r *http.Request) {
	if !g.validateApiKey(r) {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		g.logger.Warn("ws auth failed", "remote", r.RemoteAddr)
		return
	}

	conn, err := g.upgrader.Upgrade(w, r, nil)
	if err != nil {
		g.logger.Error("ws upgrade error", err, "remote", r.RemoteAddr)
		return
	}

	// Configure read limits & pong handler
	if g.bodySizeBytes > 0 {
		conn.SetReadLimit(g.bodySizeBytes)
	} else {
		conn.SetReadLimit(1 << 20) // sensible default 1MiB
	}

	_ = conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	conn.SetPongHandler(func(appData string) error {
		_ = conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})

	// First message must be a handshake identifying the connection type
	typ, data, err := conn.ReadMessage()
	if err != nil {
		g.logger.Error("failed to read handshake", "err", err, "remote", conn.RemoteAddr().String())
		_ = conn.Close()
		return
	}

	if typ != websocket.TextMessage && typ != websocket.BinaryMessage {
		g.logger.Warn("invalid handshake message type", "remote", conn.RemoteAddr().String())
		_ = writeJSONSafe(conn, map[string]string{"error": "first message must be handshake"})
		_ = conn.Close()
		return
	}

	var envelope MessageEnvelope
	if err := json.Unmarshal(data, &envelope); err != nil {
		g.logger.Warn("invalid handshake json", "err", err, "remote", conn.RemoteAddr().String())
		_ = writeJSONSafe(conn, map[string]string{"error": "invalid handshake: " + err.Error()})
		_ = conn.Close()
		return
	}

	meta := connMetadata{connType: envelope.Type}

	// Validate handshake based on type
	switch envelope.Type {
	case "mod":
		var hs ModHandshake
		if err := json.Unmarshal(envelope.Data, &hs); err != nil {
			g.logger.Warn("invalid mod handshake", "err", err, "remote", conn.RemoteAddr().String())
			_ = writeJSONSafe(conn, map[string]string{"error": "invalid mod handshake: " + err.Error()})
			_ = conn.Close()
			return
		}
		meta.id = hs.ServerID
		g.registerConn(conn, meta)
		g.logger.Info("mod connected", "server_id", hs.ServerID, "remote", conn.RemoteAddr().String())
		go g.modReadLoop(conn, meta)

	case "bot":
		var hs BotHandshake
		if err := json.Unmarshal(envelope.Data, &hs); err != nil {
			g.logger.Warn("invalid bot handshake", "err", err, "remote", conn.RemoteAddr().String())
			_ = writeJSONSafe(conn, map[string]string{"error": "invalid bot handshake: " + err.Error()})
			_ = conn.Close()
			return
		}
		meta.id = hs.BotID
		g.registerConn(conn, meta)
		g.logger.Info("bot connected", "bot_id", hs.BotID, "remote", conn.RemoteAddr().String())
		go g.botReadLoop(conn, meta)

	default:
		g.logger.Warn("unknown connection type", "type", envelope.Type, "remote", conn.RemoteAddr().String())
		_ = writeJSONSafe(conn, map[string]string{"error": "unknown connection type: " + envelope.Type})
		_ = conn.Close()
		return
	}
}

func (g *WebsocketGateway) modReadLoop(c *websocket.Conn, meta connMetadata) {
	defer func() {
		g.unregisterConn(c)
		_ = c.Close()
		g.logger.Info("mod disconnected", "server_id", meta.id, "remote", c.RemoteAddr().String())
	}()

	pingTicker := time.NewTicker(30 * time.Second)
	defer pingTicker.Stop()

	for {
		typ, data, err := c.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				g.logger.Warn("unexpected mod close", "server_id", meta.id, "err", err)
			} else {
				g.logger.Debug("mod read error", "server_id", meta.id, "err", err)
			}
			return
		}

		if typ != websocket.TextMessage && typ != websocket.BinaryMessage {
			continue
		}

		var msg GatewayModMessageIn
		if err := json.Unmarshal(data, &msg); err != nil {
			_ = writeJSONSafe(c, map[string]string{"error": "invalid json: " + err.Error()})
			g.logger.Warn("invalid json from mod", "server_id", meta.id, "remote", c.RemoteAddr().String(), "err", err)
			continue
		}

		msg.ReceivedAt = time.Now().UTC()
		if err := msg.Validate(); err != nil {
			_ = writeJSONSafe(c, map[string]string{"error": err.Error()})
			g.logger.Warn("mod message validation failed", "server_id", meta.id, "remote", c.RemoteAddr().String(), "err", err)
			continue
		}

		// Handle the message (forward to bot, enrich, etc.)
		if err := g.modHandler.Handle(context.Background(), msg); err != nil {
			_ = writeJSONSafe(c, map[string]string{"error": "handler error: " + err.Error()})
			g.logger.Error("mod handler error", "server_id", meta.id, "err", err)
			continue
		}

		_ = writeJSONSafe(c, map[string]string{"status": "ok"})

		// Handle pings
		select {
		case <-pingTicker.C:
			_ = c.SetWriteDeadline(time.Now().Add(5 * time.Second))
			if err := c.WriteMessage(websocket.PingMessage, nil); err != nil {
				g.logger.Debug("write ping failed", "server_id", meta.id, "err", err)
				return
			}
		default:
		}
	}
}

func (g *WebsocketGateway) botReadLoop(c *websocket.Conn, meta connMetadata) {
	defer func() {
		g.unregisterConn(c)
		_ = c.Close()
		g.logger.Info("bot disconnected", "bot_id", meta.id, "remote", c.RemoteAddr().String())
	}()

	pingTicker := time.NewTicker(30 * time.Second)
	defer pingTicker.Stop()

	for {
		typ, data, err := c.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				g.logger.Warn("unexpected bot close", "bot_id", meta.id, "err", err)
			} else {
				g.logger.Debug("bot read error", "bot_id", meta.id, "err", err)
			}
			return
		}

		if typ != websocket.TextMessage && typ != websocket.BinaryMessage {
			continue
		}

		var msg GatewayBotMessageIn
		if err := json.Unmarshal(data, &msg); err != nil {
			_ = writeJSONSafe(c, map[string]string{"error": "invalid json: " + err.Error()})
			g.logger.Warn("invalid json from bot", "bot_id", meta.id, "remote", c.RemoteAddr().String(), "err", err)
			continue
		}

		msg.ReceivedAt = time.Now().UTC()
		if err := msg.Validate(); err != nil {
			_ = writeJSONSafe(c, map[string]string{"error": err.Error()})
			g.logger.Warn("bot message validation failed", "bot_id", meta.id, "remote", c.RemoteAddr().String(), "err", err)
			continue
		}

		// Handle the message (forward to mod, enrich, etc.)
		if err := g.botHandler.Handle(context.Background(), msg); err != nil {
			_ = writeJSONSafe(c, map[string]string{"error": "handler error: " + err.Error()})
			g.logger.Error("bot handler error", "bot_id", meta.id, "err", err)
			continue
		}

		_ = writeJSONSafe(c, map[string]string{"status": "ok"})

		// Handle pings
		select {
		case <-pingTicker.C:
			_ = c.SetWriteDeadline(time.Now().Add(5 * time.Second))
			if err := c.WriteMessage(websocket.PingMessage, nil); err != nil {
				g.logger.Debug("write ping failed", "bot_id", meta.id, "err", err)
				return
			}
		default:
		}
	}
}