// Package server 神识 HTTP 服务：数据面（IR 入站 → SSE 出站）、健康面、用量面与 BYOK 面。
//
// 故障切换红线：首帧到达前不提交 SSE 响应——失败可整体换道并返回 JSON 错误；
// 已写出字节后不再换道，如实以 error 终帧收尾（shenshi-gateway.md §4/§5）。
//
// 配额三段式：预扣（max_tokens 上限，钳位）→ 结算（实际 token）→ 退款（失败幂等）。
// BYOK：请求携带 byok_key_id 时独占候选（不回落平台端点，防止用户预期自带密钥却被平台计费）。
package server

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"time"

	"hunjing/shenshi/internal/byok"
	"hunjing/shenshi/internal/ir"
	"hunjing/shenshi/internal/metering"
	"hunjing/shenshi/internal/quota"
	"hunjing/shenshi/internal/relay"
	"hunjing/shenshi/internal/route"
	"hunjing/shenshi/internal/store"
	"hunjing/shenshi/internal/transform"
)

// Config 服务配置。
type Config struct {
	ServiceToken string
	Store        store.Store // nil 时计量退化为进程内环形
	Vault        *byok.Vault // nil 时 BYOK 面整体 503（fail-closed）
	Quota        *quota.Limiter
}

// Server 神识 HTTP 服务。
type Server struct {
	cfg    Config
	routes *route.Registry
	usage  *metering.Log
	mux    *http.ServeMux
}

// New 构建服务并注册路由。
func New(cfg Config, routes *route.Registry) *Server {
	s := &Server{
		cfg: cfg, routes: routes,
		usage: metering.New(10_000, cfg.Store),
		mux:   http.NewServeMux(),
	}
	s.mux.HandleFunc("GET /healthz", s.handleHealth)
	s.mux.HandleFunc("GET /readyz", s.handleHealth)
	s.mux.Handle("POST /v1/chat/completions", s.auth(http.HandlerFunc(s.handleChat)))
	s.mux.Handle("GET /me/usage", s.auth(http.HandlerFunc(s.handleUsage)))
	s.mux.Handle("GET /admin/usage/sample", s.auth(http.HandlerFunc(s.handleUsageSample)))
	s.mux.Handle("GET /me/byok-keys", s.auth(http.HandlerFunc(s.handleByokList)))
	s.mux.Handle("POST /me/byok-keys", s.auth(http.HandlerFunc(s.handleByokRegister)))
	s.mux.Handle("DELETE /me/byok-keys/{id}", s.auth(http.HandlerFunc(s.handleByokRevoke)))
	return s
}

// ServeHTTP 实现 http.Handler。
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { s.mux.ServeHTTP(w, r) }

// auth 服务令牌校验中间件；未配置令牌时放行（本地开发档）。
func (s *Server) auth(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if s.cfg.ServiceToken == "" {
			next.ServeHTTP(w, r)
			return
		}
		if r.Header.Get("Authorization") != "Bearer "+s.cfg.ServiceToken {
			writeError(w, http.StatusUnauthorized, "HJ-2003", "服务令牌无效")
			return
		}
		next.ServeHTTP(w, r)
	})
}

// defaultReserve 预扣默认上限（max_tokens 未指定时按 4K 输出预扣）。
const defaultReserve = 4096

// handleChat 数据面：配额预扣 → 候选迭代（首字超时/冷却标记/粘性/BYOK 独占）→ SSE → 结算/退款。
func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) {
	var req ir.ChatRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, http.StatusBadRequest, "HJ-2002", "IR 请求解析失败: "+err.Error())
		return
	}
	if err := req.Validate(); err != nil {
		writeError(w, http.StatusBadRequest, "HJ-2002", err.Error())
		return
	}

	// 三段式之一：预扣（按输出上限；钳位集中在 quota 包）
	reserveTokens := int64(defaultReserve)
	if req.MaxTokens != nil {
		reserveTokens = int64(*req.MaxTokens)
	}
	reserved, err := s.quotaReserve(req.UserID, reserveTokens)
	if err != nil {
		writeError(w, http.StatusPaymentRequired, "HJ-2004", err.Error())
		return
	}

	candidates, err := s.candidates(&req)
	if err != nil {
		_ = s.quotaRefund(req.UserID, reserved)
		writeError(w, http.StatusServiceUnavailable, "HJ-2001", err.Error())
		return
	}

	var lastErr error
	for i, endpoint := range candidates {
		res := s.tryEndpoint(w, r, endpoint, &req, i)
		s.usage.Record(res.Attempt)
		if res.Committed {
			// 已向客户端写出：无论成败都不能再换道
			if res.Status == "ok" {
				s.routes.Sticky.Remember(req.RunID, endpoint.Name)
				s.quotaSettle(req.UserID, reserved, res.Attempt.PromptTokens+res.Attempt.CompletionTokens)
			} else {
				// 流中失败：已发生的 token 照实结算，其余退款
				s.quotaSettle(req.UserID, reserved, res.Attempt.PromptTokens+res.Attempt.CompletionTokens)
			}
			return
		}
		// 未提交：分类冷却后换下一候选
		if res.Upstream != nil {
			if res.Upstream.StatusCode == 401 || res.Upstream.StatusCode == 403 {
				s.routes.Cooldowns.Mark(route.CredentialCooldownKey(endpoint), true)
			} else {
				s.routes.Cooldowns.Mark(route.BindingCooldownKey(endpoint), false)
			}
		} else {
			s.routes.Cooldowns.Mark(route.BindingCooldownKey(endpoint), false)
		}
		lastErr = res.Err
	}
	// 三段式之三：全部候选失败 → 全额退款
	_ = s.quotaRefund(req.UserID, reserved)
	writeError(w, http.StatusServiceUnavailable, "HJ-2001", "全部候选端点失败："+errString(lastErr))
}

// candidates 候选序列：BYOK 独占（不回落平台端点）；否则注册表候选。
func (s *Server) candidates(req *ir.ChatRequest) ([]route.Endpoint, error) {
	if req.ByokKeyID == "" {
		return s.routes.Candidates(req.Slot, req.RunID)
	}
	if s.cfg.Vault == nil {
		return nil, fmt.Errorf("HJ-2003 BYOK 功能未启用（SHENSHI_BYOK_MASTER_KEY 未配置）")
	}
	key, plain, err := s.cfg.Vault.Resolve(req.UserID, req.ByokKeyID)
	if err != nil {
		return nil, err
	}
	spec := route.EndpointSpec{Type: key.Type, Name: "byok:" + key.ID, BaseURL: key.BaseURL, APIKey: plain, Model: key.Model}
	adaptor, ok := route.BuildAdaptor(spec)
	if !ok {
		return nil, fmt.Errorf("HJ-2003 BYOK 密钥类型不受支持: %s", key.Type)
	}
	return []route.Endpoint{{
		Name: spec.Name, Type: key.Type, Adaptor: adaptor,
		Priority: -100, CredentialKey: spec.Name, UpstreamModel: key.Model,
	}}, nil
}

func (s *Server) quotaReserve(userID string, tokens int64) (int64, error) {
	if s.cfg.Quota == nil {
		return 0, nil
	}
	return s.cfg.Quota.Reserve(userID, tokens)
}

func (s *Server) quotaSettle(userID string, reserved int64, actual int) error {
	if s.cfg.Quota == nil {
		return nil
	}
	return s.cfg.Quota.Settle(userID, reserved, int64(actual))
}

func (s *Server) quotaRefund(userID string, reserved int64) error {
	if s.cfg.Quota == nil {
		return nil
	}
	return s.cfg.Quota.Refund(userID, reserved)
}

type attempt struct {
	metering.Attempt
	Committed bool                     // 是否已向客户端写出 SSE
	Upstream  *transform.UpstreamError // 上游分类（可空）
	Err       error                    // 原始错误
}

func errString(err error) string {
	if err == nil {
		return "未知错误"
	}
	return err.Error()
}

// tryEndpoint 单端点尝试。
func (s *Server) tryEndpoint(w http.ResponseWriter, r *http.Request, endpoint route.Endpoint, req *ir.ChatRequest, seq int) attempt {
	start := time.Now()
	ctx, cancel := context.WithTimeout(r.Context(), s.routes.FirstByteTimeout)
	defer cancel()

	deltas := make(chan ir.Delta, 16)
	errCh := make(chan error, 1)
	go func() {
		defer close(deltas)
		errCh <- endpoint.Adaptor.Stream(ctx, req, deltas)
	}()

	res := attempt{Attempt: metering.Attempt{
		ID:    fmt.Sprintf("%s#%s#%d", req.RunID, endpoint.Name, seq),
		RunID: req.RunID, CharacterID: req.CharacterID, UserID: req.UserID,
		Model: req.Model, Slot: req.Slot, Endpoint: endpoint.Name,
	}}

	// 阶段一：等首帧或失败——此阶段未提交响应，可整体回退
	var first ir.Delta
	select {
	case d, ok := <-deltas:
		if !ok {
			res.Err = <-errCh
			res.Upstream, _ = res.Err.(*transform.UpstreamError)
			res.Status, res.TotalMS = failStatus(res.Err), time.Since(start).Milliseconds()
			return res
		}
		first = d
	case <-ctx.Done():
		res.Status, res.Err = "first_byte_timeout", ctx.Err()
		res.TotalMS = time.Since(start).Milliseconds()
		return res
	}

	// 阶段二：首帧已到——提交 SSE，流式转发
	relay.SSEHeaders(w)
	w.WriteHeader(http.StatusOK)
	res.Committed = true
	res.FirstTokenMS = time.Since(start).Milliseconds()
	_ = relay.WriteDelta(w, first)

	for d := range deltas {
		if d.Delta != "" || d.FinishReason != "" {
			if err := relay.WriteDelta(w, d); err != nil {
				break // 客户端断开
			}
		}
		if d.FinishReason != "" {
			res.Attempt.PromptTokens = d.UsagePromptTokens
			res.Attempt.CompletionTokens = d.UsageCompletionTokens
		}
	}
	err := <-errCh
	res.TotalMS = time.Since(start).Milliseconds()
	if err == nil {
		_ = relay.WriteDone(w)
		res.Status = "ok"
		return res
	}
	// 流中失败：如实以 error 终帧收尾，不伪装完成
	res.Err = err
	res.Upstream, _ = err.(*transform.UpstreamError)
	res.Status = "error"
	res.ErrorCode = "stream_written"
	_ = relay.WriteDelta(w, ir.Delta{FinishReason: "error"})
	_ = relay.WriteDone(w)
	return res
}

func failStatus(err error) string {
	if up, ok := err.(*transform.UpstreamError); ok {
		if up.StatusCode == 0 {
			return "timeout_or_network"
		}
		return "upstream_error"
	}
	return "error"
}

// meUser 从代理头取用户标识（灵台鉴权后注入，服务令牌即信任边界）。
func meUser(r *http.Request) string { return r.Header.Get("X-Hunjing-User-Id") }

// handleUsage 用户用量聚合（/me/usage；附配额余额）。
func (s *Server) handleUsage(w http.ResponseWriter, r *http.Request) {
	userID := meUser(r)
	rows, err := s.usage.Aggregate(userID)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "HJ-2001", err.Error())
		return
	}
	resp := map[string]any{"rows": rows}
	if s.cfg.Quota != nil {
		budget, used, reserved, err := s.cfg.Quota.Budget(userID)
		if err == nil {
			resp["quota"] = map[string]int64{"budget": budget, "used": used, "reserved": reserved}
		}
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(resp)
}

// handleUsageSample 对账抽样（/admin/usage/sample?limit=N；最近 N 条原始记录）。
func (s *Server) handleUsageSample(w http.ResponseWriter, r *http.Request) {
	limit := 50
	if v := r.URL.Query().Get("limit"); v != "" {
		if _, err := fmt.Sscanf(v, "%d", &limit); err != nil || limit <= 0 || limit > 1000 {
			writeError(w, http.StatusBadRequest, "HJ-2002", "limit 须为 1..1000")
			return
		}
	}
	rows, err := s.usage.Sample(limit)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "HJ-2001", err.Error())
		return
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(map[string]any{"rows": rows})
}

func (s *Server) handleByokList(w http.ResponseWriter, r *http.Request) {
	if s.cfg.Vault == nil {
		writeError(w, http.StatusServiceUnavailable, "HJ-2003", "BYOK 功能未启用（SHENSHI_BYOK_MASTER_KEY 未配置）")
		return
	}
	keys, err := s.cfg.Vault.List(meUser(r))
	if err != nil {
		writeError(w, http.StatusInternalServerError, "HJ-2001", err.Error())
		return
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(map[string]any{"keys": keys})
}

type byokRegisterRequest struct {
	Name    string `json:"name"`
	Type    string `json:"type"` // openai_compatible / anthropic / gemini
	BaseURL string `json:"base_url"`
	APIKey  string `json:"api_key"`
	Model   string `json:"model"`
}

func (s *Server) handleByokRegister(w http.ResponseWriter, r *http.Request) {
	if s.cfg.Vault == nil {
		writeError(w, http.StatusServiceUnavailable, "HJ-2003", "BYOK 功能未启用（SHENSHI_BYOK_MASTER_KEY 未配置）")
		return
	}
	var body byokRegisterRequest
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		writeError(w, http.StatusBadRequest, "HJ-2002", "请求解析失败: "+err.Error())
		return
	}
	switch body.Type {
	case "openai_compatible", "anthropic", "gemini":
	default:
		writeError(w, http.StatusBadRequest, "HJ-2002", "type 须为 openai_compatible / anthropic / gemini")
		return
	}
	if body.APIKey == "" || body.BaseURL == "" || body.Model == "" {
		writeError(w, http.StatusBadRequest, "HJ-2002", "api_key / base_url / model 不能为空")
		return
	}
	key, err := s.cfg.Vault.Register(meUser(r), body.Name, body.Type, body.BaseURL, body.APIKey, body.Model)
	if err != nil {
		writeError(w, http.StatusBadRequest, "HJ-2003", err.Error())
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusCreated)
	_ = json.NewEncoder(w).Encode(key)
}

func (s *Server) handleByokRevoke(w http.ResponseWriter, r *http.Request) {
	if s.cfg.Vault == nil {
		writeError(w, http.StatusServiceUnavailable, "HJ-2003", "BYOK 功能未启用（SHENSHI_BYOK_MASTER_KEY 未配置）")
		return
	}
	if err := s.cfg.Vault.Revoke(meUser(r), r.PathValue("id")); err != nil {
		if errors.Is(err, store.ErrByokNotFound) {
			writeError(w, http.StatusNotFound, "HJ-2003", err.Error())
			return
		}
		writeError(w, http.StatusInternalServerError, "HJ-2001", err.Error())
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

// handleHealth 健康面。
func (s *Server) handleHealth(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write([]byte(`{"status":"ok"}`))
}

// writeError 统一错误结构（API 契约 §5）。
func writeError(w http.ResponseWriter, status int, code, message string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	body := map[string]any{
		"error": map[string]string{"code": code, "message": message},
	}
	_ = json.NewEncoder(w).Encode(body)
}