package httpserver

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"reflect"
	"strings"
	"sync"
	"time"

	"git.yuntuo.hk/k2safe/orbit/internal/fault"
	"git.yuntuo.hk/k2safe/orbit/internal/jsonvalue"
	core "git.yuntuo.hk/k2safe/orbit/model"
	"git.yuntuo.hk/k2safe/orbit/platform/auth"
	"git.yuntuo.hk/k2safe/orbit/pulse/worker"
	"github.com/coder/websocket"
)

const WorkerSubprotocol = "pulse.worker.v1"

// WorkerHub 管理节点内连接生命周期；任务归属和容量仍由 MySQL 决定。
type WorkerHub struct {
	service *worker.Service
	ctx     context.Context
	cancel  context.CancelFunc
	slots   chan struct{}
	mu      sync.Mutex
	closing bool
	wg      sync.WaitGroup
	poll    time.Duration
}

func NewWorkerHub(parent context.Context, service *worker.Service) *WorkerHub {
	ctx, cancel := context.WithCancel(parent)
	return &WorkerHub{service: service, ctx: ctx, cancel: cancel, slots: make(chan struct{}, 128), poll: time.Second}
}
func RegisterWorkers(server *http.Server, hub *WorkerHub) {
	server.Handler.(*http.ServeMux).Handle("GET /v1/worker/connect", hub)
}

// Shutdown 停止派发并关闭本节点连接；不清除持久化租约，由重連核对或到期恢复处理。
func (h *WorkerHub) Shutdown(ctx context.Context) error {
	h.mu.Lock()
	h.closing = true
	h.cancel()
	h.mu.Unlock()
	done := make(chan struct{})
	go func() { h.wg.Wait(); close(done) }()
	select {
	case <-done:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

type workerFrame struct {
	Version   string          `json:"protocol_version"`
	Type      string          `json:"type"`
	MessageID core.ID         `json:"message_id"`
	ReplyTo   core.ID         `json:"reply_to,omitempty"`
	SessionID core.ID         `json:"session_id,omitempty"`
	Body      json.RawMessage `json:"body,omitempty"`
	Error     *fault.Error    `json:"error,omitempty"`
}
type outboundFrame struct {
	data []byte
	done chan error
}

func (h *WorkerHub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	requestID, _ := core.NewID()
	w.Header().Set("X-Request-ID", string(requestID))
	w.Header().Set("Cache-Control", "no-store")
	if r.URL.RawQuery != "" {
		writeError(w, string(requestID), fault.Invalid("Worker credentials and options must not be placed in URL"))
		return
	}
	values := r.Header.Values("Authorization")
	if len(values) != 1 {
		writeError(w, string(requestID), fault.Unauthorized)
		return
	}
	scheme, value, ok := strings.Cut(values[0], " ")
	if !ok || !strings.EqualFold(scheme, "Bearer") {
		writeError(w, string(requestID), fault.Unauthorized)
		return
	}
	token, err := auth.Parse(value)
	if err != nil {
		writeError(w, string(requestID), err)
		return
	}
	supported := false
	for _, header := range r.Header.Values("Sec-WebSocket-Protocol") {
		for _, protocol := range strings.Split(header, ",") {
			supported = supported || strings.TrimSpace(protocol) == WorkerSubprotocol
		}
	}
	if !supported {
		writeJSON(w, http.StatusUpgradeRequired, map[string]any{"request_id": requestID, "error": fault.New("PROTOCOL_REQUIRED", "pulse.worker.v1 subprotocol is required")})
		return
	}
	h.mu.Lock()
	if h.closing || h.ctx.Err() != nil {
		h.mu.Unlock()
		writeError(w, string(requestID), fault.Unavailable)
		return
	}
	select {
	case h.slots <- struct{}{}:
		h.wg.Add(1)
	default:
		h.mu.Unlock()
		writeJSON(w, 429, map[string]any{"request_id": requestID, "error": fault.New("CAPACITY_EXCEEDED", "node connection capacity exceeded")})
		return
	}
	h.mu.Unlock()
	defer func() { <-h.slots; h.wg.Done() }()
	authorize, cancel := context.WithTimeout(r.Context(), 5*time.Second)
	err = h.service.Authorize(authorize, token)
	cancel()
	if err != nil {
		writeError(w, string(requestID), err)
		return
	}
	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{Subprotocols: []string{WorkerSubprotocol}, CompressionMode: websocket.CompressionDisabled})
	if err != nil {
		return
	}
	defer conn.CloseNow()
	conn.SetReadLimit(512 << 10)
	ctx, cancel := context.WithCancel(h.ctx)
	defer cancel()
	helloCtx, helloCancel := context.WithTimeout(ctx, 5*time.Second)
	kind, raw, err := conn.Read(helloCtx)
	helloCancel()
	if err != nil {
		return
	}
	var hello workerFrame
	var input worker.ConnectInput
	if kind != websocket.MessageText || decodeWorkerJSON(raw, &hello) != nil || hello.Version != "1" || hello.Type != "worker.hello" || hello.MessageID.Validate() != nil || hello.SessionID != "" || hello.ReplyTo != "" || hello.Error != nil || decodeWorkerJSON(hello.Body, &input) != nil {
		_ = conn.Close(websocket.StatusPolicyViolation, "invalid worker hello")
		return
	}
	call, cancelCall := context.WithTimeout(ctx, 8*time.Second)
	binding, welcome, err := h.service.Connect(call, token, input)
	cancelCall()
	if err != nil {
		_ = writeWorkerFrame(ctx, conn, workerFrame{Version: "1", Type: "error", MessageID: requestID, ReplyTo: hello.MessageID, Error: publicWorkerError(err)})
		return
	}
	welcomeRaw, err := json.Marshal(welcome)
	if err != nil {
		return
	}
	if err := writeWorkerFrame(ctx, conn, workerFrame{Version: "1", Type: "worker.welcome", MessageID: requestID, ReplyTo: hello.MessageID, SessionID: binding.SessionID, Body: welcomeRaw}); err != nil {
		return
	}
	// 有界读写队列与 SQL 控制循环分开，处理器绝不运行在这里。
	incoming := make(chan workerFrame, 4)
	outgoing := make(chan outboundFrame, 8)
	var loops sync.WaitGroup
	loops.Add(2)
	go func() {
		defer loops.Done()
		defer cancel()
		for {
			readCtx, stop := context.WithTimeout(ctx, 60*time.Second)
			kind, raw, err := conn.Read(readCtx)
			stop()
			if err != nil {
				return
			}
			var frame workerFrame
			if kind != websocket.MessageText || decodeWorkerJSON(raw, &frame) != nil || frame.Version != "1" || frame.MessageID.Validate() != nil || frame.SessionID != binding.SessionID || frame.ReplyTo != "" || frame.Error != nil {
				_ = conn.Close(websocket.StatusPolicyViolation, "invalid worker frame")
				return
			}
			select {
			case incoming <- frame:
			case <-ctx.Done():
				return
			}
		}
	}()
	go func() {
		defer loops.Done()
		defer cancel()
		for {
			select {
			case <-ctx.Done():
				return
			case frame := <-outgoing:
				writeCtx, stop := context.WithTimeout(ctx, 5*time.Second)
				err := conn.Write(writeCtx, websocket.MessageText, frame.data)
				stop()
				if frame.done != nil {
					frame.done <- err
				}
				if err != nil {
					return
				}
			}
		}
	}()
	defer func() { cancel(); _ = conn.CloseNow(); loops.Wait() }()
	send := func(frame workerFrame, wait bool) bool {
		frame.Version = "1"
		frame.SessionID = binding.SessionID
		if frame.MessageID == "" {
			var err error
			frame.MessageID, err = core.NewID()
			if err != nil {
				return false
			}
		}
		raw, err := json.Marshal(frame)
		if err != nil {
			return false
		}
		out := outboundFrame{data: raw}
		if wait {
			out.done = make(chan error, 1)
		}
		select {
		case outgoing <- out:
		default:
			return false
		}
		if !wait {
			return true
		}
		select {
		case err := <-out.done:
			return err == nil
		case <-ctx.Done():
			return false
		}
	}
	capacity := input.Capacity
	draining := welcome.Draining
	eligible := false
	leases := map[core.ID]time.Time{}
	ticker := time.NewTicker(h.poll)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case frame := <-incoming:
			call, stop := context.WithTimeout(ctx, 8*time.Second)
			var body any
			var commandErr error
			switch frame.Type {
			case "worker.capacity":
				var in worker.CapacityInput
				if commandErr = decodeWorkerJSON(frame.Body, &in); commandErr == nil {
					var out worker.Welcome
					out, commandErr = h.service.Capacity(call, token, binding, in)
					body = out
					if commandErr == nil {
						capacity = in.Capacity
						draining = out.Draining
						eligible = true
					}
				}
			case "worker.resume":
				var ref worker.AttemptRef
				if commandErr = decodeWorkerJSON(frame.Body, &ref); commandErr == nil {
					var out worker.LeaseStatus
					out, commandErr = h.service.Resume(call, token, binding, ref)
					body = out
					if commandErr == nil {
						if out.Active {
							leases[ref.AttemptID] = out.LeaseUntil
						} else {
							delete(leases, ref.AttemptID)
						}
					}
				}
			case "task.start", "task.renew":
				var ref worker.AttemptRef
				if commandErr = decodeWorkerJSON(frame.Body, &ref); commandErr == nil {
					var out worker.CommandResult
					if frame.Type == "task.start" {
						out, commandErr = h.service.Start(call, token, binding, frame.MessageID, ref)
					} else {
						out, commandErr = h.service.Renew(call, token, binding, frame.MessageID, ref)
					}
					body = out
					if commandErr == nil {
						if out.Lease.Active {
							leases[ref.AttemptID] = out.Lease.LeaseUntil
						} else {
							delete(leases, ref.AttemptID)
						}
					}
				}
			case "task.finish":
				var in worker.FinishInput
				if commandErr = decodeWorkerJSON(frame.Body, &in); commandErr == nil {
					var out worker.CommandResult
					out, commandErr = h.service.Finish(call, token, binding, frame.MessageID, in)
					body = out
					if commandErr == nil && !out.Lease.Active {
						delete(leases, in.Ref.AttemptID)
					}
				}
			default:
				commandErr = fault.Invalid("unsupported worker message type")
			}
			stop()
			reply := workerFrame{Type: "ack", ReplyTo: frame.MessageID}
			if commandErr != nil {
				reply.Type = "error"
				reply.Error = publicWorkerError(commandErr)
			} else {
				reply.Body, commandErr = json.Marshal(body)
				if commandErr != nil {
					return
				}
			}
			fatal := commandErr != nil && (fault.Code(commandErr) == "STALE_SESSION" || fault.Code(commandErr) == "UNAUTHENTICATED" || fault.Code(commandErr) == "PERMISSION_DENIED")
			if !send(reply, fatal) || fatal {
				return
			}
		case <-ticker.C:
			// 本地记录仅减少空轮询；SQL 仍原子验证总容量、队列额度和有效会话。
			now := time.Now()
			for attempt, until := range leases {
				if !until.After(now) {
					delete(leases, attempt)
				}
			}
			if !eligible || draining || capacity == 0 || len(leases) >= int(capacity) || len(outgoing) >= 4 {
				continue
			}
			for burst := 0; burst < 4 && len(leases) < int(capacity) && len(outgoing) < 4; burst++ {
				call, stop := context.WithTimeout(ctx, 8*time.Second)
				assignment, err := h.service.Claim(call, token, binding)
				stop()
				if err != nil {
					if fault.Code(err) == "STALE_SESSION" || fault.Code(err) == "UNAUTHENTICATED" || fault.Code(err) == "PERMISSION_DENIED" {
						send(workerFrame{Type: "error", Error: publicWorkerError(err)}, true)
						return
					}
					break
				}
				if assignment == nil {
					break
				}
				raw, err := json.Marshal(assignment)
				if err != nil {
					return
				}
				leases[assignment.AttemptID] = assignment.LeaseUntil
				if !send(workerFrame{Type: "task.assign", Body: raw}, false) {
					return
				}
			}
		}
	}
}
func publicWorkerError(err error) *fault.Error {
	var known *fault.Error
	if errors.As(err, &known) {
		return known
	}
	return fault.Unavailable
}
func writeWorkerFrame(ctx context.Context, conn *websocket.Conn, frame workerFrame) error {
	raw, err := json.Marshal(frame)
	if err != nil {
		return err
	}
	bounded, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()
	return conn.Write(bounded, websocket.MessageText, raw)
}
func decodeWorkerJSON(raw []byte, out any) error {
	if len(raw) == 0 || len(raw) > 512<<10 {
		return fault.Invalid("invalid worker message size")
	}
	if _, err := jsonvalue.Canonical(raw); err != nil {
		return fault.Invalid("invalid worker JSON")
	}
	if err := exactFields(raw, reflect.TypeOf(out).Elem()); err != nil {
		return fault.Invalid("unknown or invalid worker field")
	}
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(out); err != nil {
		return fault.Invalid("invalid worker field value")
	}
	return nil
}