package ws

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/url"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"github.com/hexahost/gamecloud/node-agent/internal/config"
	"github.com/hexahost/gamecloud/node-agent/internal/protocol"
	runtimepkg "github.com/hexahost/gamecloud/node-agent/internal/runtime"
)

const (
	writeWait         = 10 * time.Second
	pongWait          = 60 * time.Second
	pingPeriod        = (pongWait * 9) / 10
	maxReconnectDelay = 60 * time.Second
)

// Client maintains an outbound WebSocket connection to the control plane.
type Client struct {
	cfg          *config.Config
	log          *slog.Logger
	runtime      *runtimepkg.Manager
	agentVersion string

	writeMu sync.Mutex
	conn    *websocket.Conn
}

// NewClient creates a WebSocket client for the control plane.
func NewClient(cfg *config.Config, log *slog.Logger, mgr *runtimepkg.Manager, agentVersion string) *Client {
	return &Client{
		cfg:          cfg,
		log:          log,
		runtime:      mgr,
		agentVersion: agentVersion,
	}
}

// SetRuntime attaches the runtime manager after construction to break init cycles.
func (c *Client) SetRuntime(mgr *runtimepkg.Manager) {
	c.runtime = mgr
}

// Run connects to the control plane and processes messages until ctx is cancelled.
func (c *Client) Run(ctx context.Context) error {
	backoff := c.cfg.ReconnectBackoff
	if backoff <= 0 {
		backoff = 5 * time.Second
	}

	for {
		if ctx.Err() != nil {
			return ctx.Err()
		}

		err := c.session(ctx)
		if ctx.Err() != nil {
			return ctx.Err()
		}

		c.log.Warn("websocket disconnected", "error", err, "retry_in", backoff.String())
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(backoff):
		}

		backoff *= 2
		if backoff > maxReconnectDelay {
			backoff = maxReconnectDelay
		}
	}
}

func (c *Client) session(ctx context.Context) error {
	conn, err := c.dial(ctx)
	if err != nil {
		return err
	}

	c.writeMu.Lock()
	c.conn = conn
	c.writeMu.Unlock()

	defer func() {
		c.writeMu.Lock()
		c.conn = nil
		c.writeMu.Unlock()
		_ = conn.Close()
	}()

	conn.SetReadDeadline(time.Now().Add(pongWait))
	conn.SetPongHandler(func(string) error {
		return conn.SetReadDeadline(time.Now().Add(pongWait))
	})

	if err := c.sendHello(ctx); err != nil {
		return fmt.Errorf("send hello: %w", err)
	}

	sessionCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	errCh := make(chan error, 2)
	go func() {
		errCh <- c.heartbeatLoop(sessionCtx)
	}()
	go func() {
		errCh <- c.readLoop(sessionCtx)
	}()

	select {
	case <-ctx.Done():
		cancel()
		return ctx.Err()
	case err := <-errCh:
		cancel()
		return err
	}
}

func (c *Client) dial(ctx context.Context) (*websocket.Conn, error) {
	endpoint, err := c.buildURL()
	if err != nil {
		return nil, err
	}

	dialer := websocket.DefaultDialer
	conn, _, err := dialer.DialContext(ctx, endpoint, nil)
	if err != nil {
		return nil, fmt.Errorf("dial websocket: %w", err)
	}
	c.log.Info("connected to control plane", "url", redactToken(endpoint))
	return conn, nil
}

func (c *Client) buildURL() (string, error) {
	u, err := url.Parse(c.cfg.ControlPlaneURL)
	if err != nil {
		return "", fmt.Errorf("parse control plane url: %w", err)
	}
	q := u.Query()
	q.Set("nodeId", c.cfg.NodeID)
	if c.cfg.NodeToken != "" {
		q.Set("token", c.cfg.NodeToken)
	}
	u.RawQuery = q.Encode()
	return u.String(), nil
}

func (c *Client) sendHello(ctx context.Context) error {
	hostname, _ := os.Hostname()
	payload := protocol.AgentHelloPayload{
		NodeID:          c.cfg.NodeID,
		AgentVersion:    c.agentVersion,
		ProtocolVersion: protocol.CurrentProtocolVersion,
		Hostname:        hostname,
		Capabilities:    []string{"minecraft", "docker"},
	}
	return c.send(ctx, protocol.TypeAgentHello, "", payload)
}

func (c *Client) heartbeatLoop(ctx context.Context) error {
	ticker := time.NewTicker(c.cfg.HeartbeatInterval)
	defer ticker.Stop()

	pingTicker := time.NewTicker(pingPeriod)
	defer pingTicker.Stop()

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			if err := c.sendHeartbeat(ctx); err != nil {
				return err
			}
		case <-pingTicker.C:
			if err := c.writePing(); err != nil {
				return err
			}
		}
	}
}

func (c *Client) sendHeartbeat(ctx context.Context) error {
	var memStats runtime.MemStats
	runtime.ReadMemStats(&memStats)

	payload := protocol.AgentHeartbeatPayload{
		NodeID:           c.cfg.NodeID,
		CPUUsagePercent:  0,
		MemoryUsedBytes:  int64(memStats.Alloc),
		MemoryTotalBytes: int64(memStats.Sys),
		RunningServers:   c.runtime.RunningCount(),
	}
	return c.send(ctx, protocol.TypeAgentHeartbeat, "", payload)
}

func (c *Client) readLoop(ctx context.Context) error {
	for {
		if ctx.Err() != nil {
			return ctx.Err()
		}

		c.writeMu.Lock()
		conn := c.conn
		c.writeMu.Unlock()
		if conn == nil {
			return fmt.Errorf("connection closed")
		}

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

		env, err := protocol.DecodeEnvelope(data)
		if err != nil {
			c.log.Warn("invalid inbound envelope", "error", err)
			continue
		}

		c.dispatch(ctx, env)
	}
}

func (c *Client) dispatch(ctx context.Context, env *protocol.Envelope) {
	correlationID := env.CorrelationID
	if correlationID == "" {
		correlationID = env.MessageID
	}

	switch env.Type {
	case protocol.TypeServerProvision:
		var payload protocol.ServerProvisionPayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.provision", "error", err)
			return
		}
		c.runtime.HandleProvision(ctx, correlationID, payload)

	case protocol.TypeServerStart:
		var payload protocol.ServerStartPayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.start", "error", err)
			return
		}
		c.runtime.HandleStart(ctx, correlationID, payload)

	case protocol.TypeServerStop:
		var payload protocol.ServerStopPayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.stop", "error", err)
			return
		}
		c.runtime.HandleStop(ctx, correlationID, payload)

	case protocol.TypeServerDelete:
		var meta protocol.OperationMeta
		if err := json.Unmarshal(env.Payload, &meta); err != nil {
			c.log.Warn("decode server.delete", "error", err)
			return
		}
		c.runtime.HandleDelete(ctx, correlationID, meta)

	case protocol.TypeServerLogsSubscribe:
		var payload protocol.ServerLogsSubscribePayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.logs.subscribe", "error", err)
			return
		}
		c.runtime.HandleLogsSubscribe(ctx, env.MessageID, payload)

	case protocol.TypeServerCommand:
		var payload protocol.ServerCommandPayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.command", "error", err)
			return
		}
		c.runtime.HandleCommand(ctx, env.MessageID, payload)

	case protocol.TypeServerFilesList:
		var payload protocol.ServerFilesListPayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.files.list", "error", err)
			return
		}
		c.runtime.HandleFilesList(ctx, env.MessageID, payload)

	case protocol.TypeServerFilesRead:
		var payload protocol.ServerFilesReadPayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.files.read", "error", err)
			return
		}
		c.runtime.HandleFilesRead(ctx, env.MessageID, payload)

	case protocol.TypeServerFilesWrite:
		var payload protocol.ServerFilesWritePayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.files.write", "error", err)
			return
		}
		c.runtime.HandleFilesWrite(ctx, env.MessageID, payload)

	case protocol.TypeServerBackupPrepare:
		var meta protocol.OperationMeta
		if err := json.Unmarshal(env.Payload, &meta); err != nil {
			c.log.Warn("decode server.backup.prepare", "error", err)
			return
		}
		c.runtime.HandleBackupPrepare(ctx, env.MessageID, meta)

	case protocol.TypeServerBackupRelease:
		var meta protocol.OperationMeta
		if err := json.Unmarshal(env.Payload, &meta); err != nil {
			c.log.Warn("decode server.backup.release", "error", err)
			return
		}
		c.runtime.HandleBackupRelease(ctx, env.MessageID, meta)

	case protocol.TypeServerWorldValidate:
		var payload protocol.ServerWorldValidatePayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.world.validate", "error", err)
			return
		}
		c.runtime.HandleWorldValidate(ctx, env.MessageID, payload)

	case protocol.TypeServerWorldArchive:
		var payload protocol.ServerWorldArchivePayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.world.archive", "error", err)
			return
		}
		c.runtime.HandleWorldArchive(ctx, env.MessageID, payload)

	case protocol.TypeServerStorageUpload:
		var payload protocol.ServerStorageUploadPayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.storage.upload", "error", err)
			return
		}
		c.runtime.HandleStorageUpload(ctx, env.MessageID, payload)

	case protocol.TypeServerStorageDownload:
		var payload protocol.ServerStorageDownloadPayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.storage.download", "error", err)
			return
		}
		c.runtime.HandleStorageDownload(ctx, env.MessageID, payload)

	case protocol.TypeServerWorldReplace:
		var payload protocol.ServerWorldReplacePayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.world.replace", "error", err)
			return
		}
		c.runtime.HandleWorldReplace(ctx, env.MessageID, payload)

	case protocol.TypeServerAddonInstall:
		var payload protocol.ServerAddonInstallPayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.addon.install", "error", err)
			return
		}
		c.runtime.HandleAddonInstall(ctx, env.MessageID, payload)

	case protocol.TypeServerAddonRemove:
		var payload protocol.ServerAddonRemovePayload
		if err := json.Unmarshal(env.Payload, &payload); err != nil {
			c.log.Warn("decode server.addon.remove", "error", err)
			return
		}
		c.runtime.HandleAddonRemove(ctx, env.MessageID, payload)

	default:
		c.log.Debug("ignored control message", "type", env.Type)
	}
}

func (c *Client) send(ctx context.Context, msgType protocol.MessageType, correlationID string, payload any) error {
	env, err := protocol.NewEnvelope(msgType, uuid.NewString(), payload)
	if err != nil {
		return err
	}
	if correlationID != "" {
		env.CorrelationID = correlationID
	}
	return c.writeEnvelope(ctx, env)
}

func (c *Client) writeEnvelope(ctx context.Context, env *protocol.Envelope) error {
	data, err := json.Marshal(env)
	if err != nil {
		return fmt.Errorf("marshal envelope: %w", err)
	}

	c.writeMu.Lock()
	conn := c.conn
	c.writeMu.Unlock()
	if conn == nil {
		return fmt.Errorf("not connected")
	}

	deadline, ok := ctx.Deadline()
	if !ok {
		deadline = time.Now().Add(writeWait)
	}
	if err := conn.SetWriteDeadline(deadline); err != nil {
		return err
	}

	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	if c.conn == nil {
		return fmt.Errorf("not connected")
	}
	return c.conn.WriteMessage(websocket.TextMessage, data)
}

func (c *Client) writePing() error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	if c.conn == nil {
		return fmt.Errorf("not connected")
	}
	return c.conn.WriteMessage(websocket.PingMessage, nil)
}

// SendState implements runtime.Responder.
func (c *Client) SendState(ctx context.Context, correlationID string, payload protocol.ServerStatePayload) error {
	return c.send(ctx, protocol.TypeServerState, correlationID, payload)
}

// SendLog implements runtime.Responder.
func (c *Client) SendLog(ctx context.Context, correlationID string, payload protocol.ServerLogPayload) error {
	return c.send(ctx, protocol.TypeServerLog, correlationID, payload)
}

// SendOperationComplete implements runtime.Responder.
func (c *Client) SendOperationComplete(ctx context.Context, correlationID string, payload protocol.ServerOperationResultPayload) error {
	return c.send(ctx, protocol.TypeServerOperationComplete, correlationID, payload)
}

// SendOperationFailed implements runtime.Responder.
func (c *Client) SendOperationFailed(ctx context.Context, correlationID string, payload protocol.ServerOperationResultPayload) error {
	return c.send(ctx, protocol.TypeServerOperationFailed, correlationID, payload)
}

func redactToken(rawURL string) string {
	u, err := url.Parse(rawURL)
	if err != nil {
		return rawURL
	}
	if u.Query().Get("token") != "" {
		q := u.Query()
		q.Set("token", "REDACTED")
		u.RawQuery = q.Encode()
	}
	return u.String()
}