package handler

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

	"github.com/gogf/gf/v2/frame/g"
	"github.com/gorilla/websocket"

	"github.com/qianfree/team-api/internal/dispatchadapter"
	"github.com/qianfree/team-api/internal/logic/monitor"
	"github.com/qianfree/team-api/relay/channel"
	"github.com/qianfree/team-api/relay/channel/openai"
	"github.com/qianfree/team-api/relay/common"
	"github.com/qianfree/team-api/relay/constant"
	"github.com/qianfree/team-api/relaykit/dispatch"
)

// websocketUpgrader HTTP → WebSocket 升级器
var websocketUpgrader = websocket.Upgrader{
	Subprotocols: []string{"realtime"},
	CheckOrigin:  checkRealtimeOrigin,
}

// checkRealtimeOrigin 校验 WebSocket 握手来源，防跨站 WebSocket 劫持（CSWSH，B15）。
// 策略：
//   - 无 Origin 头：放行。Realtime API 的主要客户端是原生 SDK/CLI/服务端，不发送 Origin；
//     浏览器才强制带 Origin，据此区分。
//   - 同源（Origin.host == 请求 Host）：放行。
//   - 命中 REALTIME_ALLOWED_ORIGINS（逗号分隔的可信来源，如租户控制台/playground 域名）：放行。
//   - 其余跨源浏览器请求：拒绝。
func checkRealtimeOrigin(r *http.Request) bool {
	origin := r.Header.Get("Origin")
	if origin == "" {
		return true // 非浏览器客户端
	}

	u, err := url.Parse(origin)
	if err != nil || u.Host == "" {
		return false
	}

	if strings.EqualFold(u.Host, r.Host) {
		return true // 同源
	}

	for _, allowed := range strings.Split(os.Getenv("REALTIME_ALLOWED_ORIGINS"), ",") {
		allowed = strings.TrimSpace(allowed)
		if allowed == "" {
			continue
		}
		if strings.EqualFold(origin, allowed) || strings.EqualFold(u.Host, allowed) {
			return true
		}
	}

	return false
}

// RealtimeContext Realtime 请求上下文
type RealtimeContext struct {
	TenantID        int64
	UserID          int64
	ApiKeyID        int64
	ProjectID       int64 // 通过 API Key 关联的项目 ID
	RequestID       string
	ClientIP        string
	Scope           string
	KeyRateLimitQps int
	KeyConcurrency  int
	KeyIpWhitelist  string
	KeyTotalQuota   float64
	KeyUsedQuota    float64
}

// HandleRealtime 处理 /v1/realtime WebSocket 请求
func HandleRealtime(w http.ResponseWriter, r *http.Request, rc *RealtimeContext, provider common.DataProvider, billing common.BillingProvider) (*common.Usage, *BillingResult, error) {
	ctx := r.Context()

	// 0. 前置检查（scope / IP 白名单 / QPS 限流）：只依赖认证上下文，放在 WebSocket
	// 升级之前——升级后连接已被劫持，无法再写标准 HTTP 状态码；前置后被拒请求能拿到
	// 正常的 401/429 响应，也省掉无谓的升级握手和首消息读取。
	// 注意：上层调用方（internal/handler/relay）对 realtime 的返回错误只记日志、不写响应
	//（因升级后的错误无法写回），故升级前的拒绝必须在此处自行写回 HTTP 错误响应，
	// 否则客户端会收到空 body 的 200。
	billingResult := &BillingResult{}
	if billing != nil {
		if !billing.CheckScope(rc.Scope, "realtime") {
			err := constant.NewAuthError("API key scope denied")
			WriteRelayError(w, err)
			return nil, billingResult, err
		}
		if !billing.CheckIPWhitelist(rc.KeyIpWhitelist, rc.ClientIP) {
			err := constant.NewAuthError("IP address is not allowed")
			WriteRelayError(w, err)
			return nil, billingResult, err
		}
		allowed, limitLevel, limit, remaining, resetAt := billing.CheckRateLimit(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, rc.KeyRateLimitQps)
		if !allowed {
			err := &RelayErrorWithRateLimit{
				StatusCode: 429,
				Message:    fmt.Sprintf("rate limit exceeded at %s level", limitLevel),
				LimitLevel: limitLevel,
				Remaining:  remaining,
				ResetAt:    resetAt,
			}
			WriteRelayError(w, err)
			return nil, billingResult, err
		}
		billingResult.RateLimitInfo = &common.RateLimitInfo{
			Limit:     limit,
			Remaining: remaining,
			ResetAt:   resetAt,
		}
	}

	// 1. WebSocket 升级
	clientConn, err := websocketUpgrader.Upgrade(w, r, nil)
	if err != nil {
		g.Log().Errorf(ctx, "[HandleRealtime] WebSocket upgrade failed: %v", err)
		return nil, nil, fmt.Errorf("websocket upgrade failed: %w", err)
	}
	defer clientConn.Close()

	// 2. 读取第一条消息获取 model 名称
	_, firstMsg, err := clientConn.ReadMessage()
	if err != nil {
		return nil, nil, fmt.Errorf("read first websocket message failed: %w", err)
	}

	var firstEvent map[string]json.RawMessage
	if err := json.Unmarshal(firstMsg, &firstEvent); err != nil {
		return nil, nil, fmt.Errorf("parse first websocket message failed: %w", err)
	}

	modelName := extractModelFromEvent(firstEvent)
	if modelName == "" {
		errMsg, _ := json.Marshal(map[string]any{
			"type":  "error",
			"error": map[string]string{"message": "model is required in the first message"},
		})
		_ = clientConn.WriteMessage(websocket.TextMessage, errMsg)
		return nil, nil, fmt.Errorf("model not found in first websocket message")
	}

	// 3. 验证模型
	_, _, err = provider.GetModelMapping(ctx, modelName)
	if err != nil {
		return nil, nil, fmt.Errorf("model not found: %s", modelName)
	}

	// 3.1 检查成员模型范围
	if allowed, err := provider.CheckMemberModelAccess(ctx, rc.TenantID, rc.UserID, modelName); err != nil {
		return nil, nil, err
	} else if !allowed {
		return nil, nil, constant.NewAuthError("model not allowed for this member")
	}
	if allowed, err := provider.CheckApiKeyModelAccess(ctx, rc.ApiKeyID, modelName); err != nil {
		return nil, nil, err
	} else if !allowed {
		return nil, nil, constant.NewAuthError("model not allowed for this API key")
	}

	// 4. 渠道选择（新调度引擎；websocket 长连接由续期器保活租约）
	enabled, channelScope, err := provider.CheckTenantModelAccess(ctx, rc.TenantID, modelName)
	if err != nil {
		return nil, nil, constant.NewChannelError("no available channel for model: "+modelName, err)
	}
	if !enabled {
		return nil, nil, constant.NewChannelError("no available channel for model: "+modelName, common.ErrTenantModelNotEnabled)
	}
	sess := dispatchadapter.Coordinator(ctx).Route(ctx, dispatch.RequestProfile{
		RequestID: rc.RequestID,
		TenantID:  rc.TenantID,
		UserID:    rc.UserID,
		APIKeyID:  rc.ApiKeyID,
		Model:     modelName,
		Scope:     channelScope,
		Replay:    dispatch.ReplayCostly,
		Signals:   dispatch.SessionSignals{HeaderSessionID: r.Header.Get("X-Session-Id")},
		Policy:    dispatchadapter.TenantRoutingPolicy(ctx, rc.TenantID),
	})
	defer sess.Finish(context.WithoutCancel(ctx), false, 0)

	var selection *common.ChannelSelection
	for {
		d := sess.Next(ctx)
		if d == nil {
			monitor.TrackDispatchNoCandidate()
			return nil, nil, constant.NewChannelError("no available channel for model: "+modelName, common.ErrChannelUnavailable)
		}
		monitor.TrackDispatchSelection(string(d.Reason), string(d.Channel.Tier), string(d.SessionKey.Source))
		sel, mErr := provider.MaterializeSelection(ctx, d.Channel.ID, d.KeyID, modelName)
		if mErr != nil {
			if decision := reportMaterializeFailure(ctx, sess, mErr); decision == dispatch.DecisionAbort {
				return nil, nil, constant.NewChannelError("no available channel for model: "+modelName, mErr)
			}
			continue
		}
		sel.SelectionReason = string(d.Reason)
		selection = sel
		break
	}
	leaseRefresh := startDispatchLeaseRefresher(selection.ChannelID, rc.RequestID)
	defer leaseRefresh.Stop()
	realtimeStart := time.Now()

	// 5. 构造 RelayInfo
	info := &common.RelayInfo{
		Context:         ctx,
		TenantID:        rc.TenantID,
		UserID:          rc.UserID,
		ApiKeyID:        rc.ApiKeyID,
		RequestID:       rc.RequestID,
		RelayMode:       int(constant.RelayModeRealtime),
		OriginModelName: modelName,
		RequestURLPath:  "/realtime",
		RequestHeaders:  r.Header,
		StartTime:       time.Now(),
		InboundFormat:   constant.RelayFormatOpenAI,
		ClientFormat:    constant.RelayFormatOpenAI,
		ChannelMeta: &common.ChannelMeta{
			ChannelID:         selection.ChannelID,
			ChannelType:       selection.ChannelType,
			ChannelName:       selection.ChannelName,
			BaseURL:           selection.BaseURL,
			ApiKey:            selection.ApiKey,
			UpstreamModelName: selection.UpstreamModelName,
			IsModelMapped:     selection.IsModelMapped,
			Settings:          selection.Settings,
		},
		ClientConn: clientConn,
	}

	// 6. 获取适配器
	adaptor := channel.GetAdaptor(selection.ChannelType)
	if adaptor == nil {
		return nil, nil, fmt.Errorf("unsupported channel type for realtime: %d", selection.ChannelType)
	}
	adaptor.Init(info)

	// 7. 预扣费用（Realtime 使用 0 input tokens 作为初始估算）
	var preDeductAmount float64
	if billing != nil {
		if !billing.AcquireConcurrent(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, modelName) {
			return nil, billingResult, constant.NewRateLimitError("concurrent request limit exceeded")
		}
		defer billing.ReleaseConcurrent(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, modelName)

		if !billing.AcquireApiKeyConcurrent(ctx, rc.ApiKeyID, rc.KeyConcurrency) {
			return nil, billingResult, constant.NewRateLimitError("API key concurrent request limit exceeded")
		}
		defer billing.ReleaseApiKeyConcurrent(ctx, rc.ApiKeyID)

		amt, billErr := billing.PreDeduct(ctx, rc.TenantID, modelName, 0, 0, false, rc.RequestID)
		if billErr != nil {
			// 未配价模型 fail-closed：返回明确的请求错误，而非误导性的"余额不足"
			if errors.Is(billErr, common.ErrModelPricingNotConfigured) {
				return nil, nil, constant.NewRequestError("model pricing not configured: "+modelName, billErr)
			}
			return nil, nil, constant.NewQuotaError("insufficient balance", billErr)
		}
		preDeductAmount = amt
		// 成员/Key 额度检查只做这一次（带实际预扣额），与 RelayHandler 一致：
		// 预扣前的 Check(0) 快速闸门对全量请求多付 Redis 往返，收益不成比例，已移除。
		if err := billing.CheckApiKeyQuota(ctx, rc.ApiKeyID, amt); err != nil {
			_ = billing.SettleFailed(ctx, rc.TenantID, rc.RequestID, amt)
			return nil, billingResult, constant.NewQuotaError("API key quota exceeded", err)
		}
		if err := billing.CheckMemberQuota(ctx, rc.TenantID, rc.UserID, amt); err != nil {
			_ = billing.SettleFailed(ctx, rc.TenantID, rc.RequestID, amt)
			return nil, billingResult, constant.NewQuotaError("member quota exceeded", err)
		}
	}

	billingResult.PreDeductAmount = preDeductAmount

	// 8. 建立 Realtime 代理
	proxy := openai.NewRealtimeProxy(info)
	if err := proxy.DialUpstream(); err != nil {
		// 拨号失败上报健康（绑定不删除，由守卫自然失效）；realtime 不做跨渠道重试，保持原语义
		_, _ = sess.Report(ctx, 0, err, deliveryStateOfRequestErr(err), float64(time.Since(realtimeStart).Milliseconds()), 0)
		if billing != nil && preDeductAmount > 0 {
			_ = billing.SettleFailed(ctx, rc.TenantID, rc.RequestID, preDeductAmount)
		}
		return nil, billingResult, err
	}
	defer proxy.Close()

	// 把第一条消息转发到上游
	if err := proxy.GetTargetConn().WriteMessage(websocket.TextMessage, firstMsg); err != nil {
		if billing != nil && preDeductAmount > 0 {
			_ = billing.SettleFailed(ctx, rc.TenantID, rc.RequestID, preDeductAmount)
		}
		return nil, billingResult, err
	}

	// 9. 启动双向代理
	usage, proxyErr := proxy.Proxy(ctx)
	if proxyErr == nil {
		// 成功：绑定续期 + 健康上报由 Finish 完成
		sess.Finish(context.WithoutCancel(ctx), true, float64(time.Since(realtimeStart).Milliseconds()))
	} else {
		// 会话中断：响应已开始，不可重试，仅上报健康
		_, _ = sess.Report(context.WithoutCancel(ctx), 0, proxyErr, dispatch.DeliveryResponseStarted,
			float64(time.Since(realtimeStart).Milliseconds()), 0)
	}

	// 10. 结算费用
	if billing != nil && preDeductAmount > 0 {
		if proxyErr != nil {
			streamUsage := usage
			if streamUsage == nil {
				streamUsage = &common.Usage{}
			}
			settleResult, _ := billing.SettleStreamInterrupted(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, selection.ChannelID,
				modelName, rc.RequestID, "realtime", streamUsage, preDeductAmount, rc.ProjectID, info.StartTime)
			if settleResult != nil && settleResult.ActualCost > 0 && !settleResult.DuplicateSkip {
				billing.IncrMemberQuotaUsed(ctx, rc.TenantID, rc.UserID, settleResult.ActualCost)
				billing.IncrApiKeyQuotaUsed(ctx, rc.ApiKeyID, settleResult.ActualCost)
			}
		} else if usage != nil {
			settleResult, _ := billing.Settle(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, selection.ChannelID,
				modelName, rc.RequestID, "realtime", usage, preDeductAmount, rc.ProjectID, info.StartTime)
			if settleResult != nil && settleResult.ActualCost > 0 && !settleResult.DuplicateSkip {
				billing.IncrMemberQuotaUsed(ctx, rc.TenantID, rc.UserID, settleResult.ActualCost)
				billing.IncrApiKeyQuotaUsed(ctx, rc.ApiKeyID, settleResult.ActualCost)
			}
		}
	}

	return usage, billingResult, proxyErr
}

// extractModelFromEvent 从 WebSocket 事件中提取模型名称
func extractModelFromEvent(event map[string]json.RawMessage) string {
	// 顶层 model 字段
	if m, ok := event["model"]; ok {
		var s string
		if json.Unmarshal(m, &s) == nil {
			return strings.Trim(s, `"`)
		}
	}

	// session.update 的 session.model
	if sessionRaw, ok := event["session"]; ok {
		var sess map[string]json.RawMessage
		if json.Unmarshal(sessionRaw, &sess) == nil {
			if m, ok := sess["model"]; ok {
				var s string
				if json.Unmarshal(m, &s) == nil {
					return strings.Trim(s, `"`)
				}
			}
		}
	}

	return ""
}