package httpserver

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

	"git.yuntuo.hk/k2safe/orbit/internal/dispatch"
	"git.yuntuo.hk/k2safe/orbit/internal/fault"
	"git.yuntuo.hk/k2safe/orbit/internal/jsonvalue"
	"git.yuntuo.hk/k2safe/orbit/internal/jsonwire"
	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
	dispatch   *dispatch.Pool
	emptyPolls map[[32]byte]time.Time
}

func NewWorkerHub(parent context.Context, service *worker.Service) *WorkerHub {
	ctx, cancel := context.WithCancel(parent)
	h := &WorkerHub{service: service, ctx: ctx, cancel: cancel, slots: make(chan struct{}, 128), poll: time.Second, emptyPolls: make(map[[32]byte]time.Time)}
	h.dispatch = dispatch.New(ctx, 4, 128)
	h.wg.Add(1)
	go func() { defer h.wg.Done(); <-h.dispatch.Done() }()
	return h
}
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
	}
	hello, decodeErr := decodeWorkerRequest(raw, true)
	var input worker.ConnectInput
	if kind != websocket.MessageText || decodeErr != 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 := jsonwire.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
			}
			frame, decodeErr := decodeWorkerRequest(raw, false)
			if kind != websocket.MessageText || decodeErr != nil || frame.SessionID != binding.SessionID {
				_ = 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 := jsonwire.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{}
	cancelSent := map[core.ID]bool{}
	results := make(chan claimResult, 1)
	pending, burstLeft := false, 0
	pollKey := workerPollKey(token.NamespaceID, input)
	canRequest := func() bool {
		return eligible && !draining && capacity > 0 && len(leases) < int(capacity) && len(outgoing) < 4
	}
	requestClaim := func() {
		if pending || burstLeft == 0 || !canRequest() {
			return
		}
		pending = h.dispatch.Submit(dispatch.Job{Context: ctx, Application: string(binding.ApplicationID), Queue: string(token.NamespaceID) + "/" + string(input.QueueID), Run: func(jobCtx context.Context) {
			results <- h.claimFair(jobCtx, token, binding, pollKey)
		}})
		if pending {
			burstLeft--
		}
	}
	ticker := time.NewTicker(h.poll)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case result := <-results:
			pending = false
			if result.err != nil {
				if fault.Code(result.err) == "STALE_SESSION" || fault.Code(result.err) == "UNAUTHENTICATED" || fault.Code(result.err) == "PERMISSION_DENIED" {
					send(workerFrame{Type: "error", Error: publicWorkerError(result.err)}, true)
					return
				}
				burstLeft = 0
				continue
			}
			if result.assignment == nil {
				burstLeft = 0
				continue
			}
			assignment := result.assignment
			raw, err := jsonwire.Marshal(assignment)
			if err != nil {
				return
			}
			leases[assignment.AttemptID] = assignment.LeaseUntil
			if !send(workerFrame{Type: "task.assign", Body: raw}, false) {
				return
			}
			requestClaim()
		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 = jsonwire.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:
			// 取消与领取容量、排空无关；通知丢失后仍可由 Resume/Renew 核对持久化意图。
			if input.CancellationNotifications && len(leases) > 0 {
				call, stop := context.WithTimeout(ctx, 8*time.Second)
				items, err := h.service.Cancellations(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
					}
				} else {
					active := make(map[core.ID]bool, len(items))
					for _, item := range items {
						active[item.AttemptID] = true
						if cancelSent[item.AttemptID] || len(outgoing) >= 4 {
							continue
						}
						raw, err := jsonwire.Marshal(item)
						if err != nil || !send(workerFrame{Type: "task.cancel", Body: raw}, false) {
							return
						}
						cancelSent[item.AttemptID] = true
					}
					for id := range cancelSent {
						if !active[id] {
							delete(cancelSent, id)
						}
					}
				}
			}
			// 本地记录仅减少空轮询；SQL 仍原子验证总容量、队列额度和有效会话。
			now := time.Now()
			for attempt, until := range leases {
				if !until.After(now) {
					delete(leases, attempt)
				}
			}
			if !pending {
				burstLeft = 4
				requestClaim()
			}
		}
	}
}
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 := jsonwire.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")
	}
	if err := requiredWorkerFields(raw, reflect.TypeOf(out).Elem()); err != nil {
		return err
	}
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(out); err != nil {
		return fault.Invalid("invalid worker field value")
	}
	return nil
}

// decodeWorkerRequest 拒绝请求中伪装为响应的字段；空字符串/null 不能绕过方向约束。
func decodeWorkerRequest(raw []byte, hello bool) (workerFrame, error) {
	var frame workerFrame
	if err := decodeWorkerJSON(raw, &frame); err != nil {
		return frame, err
	}
	if frame.Version != "1" || frame.MessageID.Validate() != nil || len(frame.Body) == 0 {
		return frame, fault.Invalid("invalid worker envelope")
	}
	var fields map[string]json.RawMessage
	_ = json.Unmarshal(raw, &fields)
	if fields["error"] != nil || fields["reply_to"] != nil {
		return frame, fault.Invalid("response fields are not allowed in worker requests")
	}
	if hello {
		if frame.Type != "worker.hello" || fields["session_id"] != nil {
			return frame, fault.Invalid("invalid worker hello envelope")
		}
	} else if frame.SessionID.Validate() != nil {
		return frame, fault.Invalid("worker session is required")
	}
	return frame, nil
}

// requiredWorkerFields 校验必填字段及规范 uint64 字符串，不把省略布尔值等同于显式控制意图。
func requiredWorkerFields(raw []byte, t reflect.Type) error {
	if t == reflect.TypeOf(json.RawMessage{}) || t == reflect.TypeOf(time.Time{}) {
		return nil
	}
	if t.Kind() == reflect.Pointer {
		return nil
	}
	if t.Kind() == reflect.Array || t.Kind() == reflect.Slice {
		var values []json.RawMessage
		if err := json.Unmarshal(raw, &values); err != nil {
			return fault.Invalid("invalid worker array")
		}
		for _, value := range values {
			if err := requiredWorkerFields(value, t.Elem()); err != nil {
				return err
			}
		}
		return nil
	}
	if t.Kind() != reflect.Struct {
		return nil
	}
	var fields map[string]json.RawMessage
	if json.Unmarshal(raw, &fields) != nil {
		return fault.Invalid("invalid worker object")
	}
	for i := 0; i < t.NumField(); i++ {
		field := t.Field(i)
		tag := strings.Split(field.Tag.Get("json"), ",")
		if tag[0] == "" || tag[0] == "-" {
			continue
		}
		value, found := fields[tag[0]]
		optional, numberString := false, false
		for _, flag := range tag[1:] {
			optional = optional || flag == "omitempty"
			numberString = numberString || flag == "string"
		}
		if !found {
			if optional {
				continue
			}
			return fault.Invalid("missing required worker field")
		}
		if field.Type.Kind() == reflect.Uint64 && numberString {
			var text string
			if json.Unmarshal(value, &text) != nil {
				return fault.Invalid("worker counter must be a decimal string")
			}
			n, err := strconv.ParseUint(text, 10, 64)
			if err != nil || n == 0 || strconv.FormatUint(n, 10) != text {
				return fault.Invalid("invalid worker counter")
			}
		}
		if err := requiredWorkerFields(value, field.Type); err != nil {
			return err
		}
	}
	return nil
}