// gummy provider：阿里云百炼 DashScope 流式语音识别（gummy-realtime-v1）。
// WS 协议：run-task → 推二进制音频 → finish-task；下行 result-generated / task-finished / task-failed。
// ⚠️ 待真实 DASHSCOPE_API_KEY 到位后联调验证（6B 验收项），协议细节以
// https://help.aliyun.com/zh/model-studio/ 实时语音识别文档为准。
package asr

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

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
)

const dashscopeWS = "wss://dashscope.aliyuncs.com/api-ws/v1/inference"

type GummyProvider struct {
	APIKey string
	Model  string // 默认 gummy-realtime-v1

	// spare 预连接（12B/#24）：后台常备一条已完成 run-task 握手的 DashScope
	// 会话，start 直接取用，把 dial+task-started（实测暖路径 ~130ms、抖动可达
	// 2.5~3.7s）从 first_partial 关键路径上移除。协议固定 16kHz，spare 可互换。
	mu          sync.Mutex
	spare       *spareEntry
	dialing     bool
	maintainOne sync.Once
}

type spareEntry struct {
	sess *gummySession
	born time.Time
}

// spareMaxAge spare 的可用年龄上限：DashScope 空闲 60s 断连（实测 60s 存活、
// 120s 报 Idle timeout），留余量在 45s 内取用、40s 定期换新。
const spareMaxAge = 45 * time.Second

func NewGummy(apiKey string) *GummyProvider {
	return &GummyProvider{APIKey: apiKey, Model: "gummy-realtime-v1"}
}

func (g *GummyProvider) Name() string { return "gummy" }

func (g *GummyProvider) StartSession(ctx context.Context, cfg SessionConfig) (Session, error) {
	g.maintainOne.Do(func() { go g.maintainSpare() })
	if cfg.SampleRate == 16000 {
		if s := g.takeSpare(); s != nil {
			go g.replenish()
			return s, nil
		}
		defer func() { go g.replenish() }()
	}
	return g.dialSession(ctx, cfg)
}

// takeSpare 取走当前 spare（若仍新鲜且存活）；过期/已死的就地关闭丢弃。
func (g *GummyProvider) takeSpare() *gummySession {
	g.mu.Lock()
	sp := g.spare
	g.spare = nil
	g.mu.Unlock()
	if sp == nil {
		return nil
	}
	if time.Since(sp.born) < spareMaxAge && sp.sess.alive() {
		slog.Info("gummy: session from spare", "task", sp.sess.taskID,
			"age_ms", time.Since(sp.born).Milliseconds())
		return sp.sess
	}
	_ = sp.sess.Close()
	return nil
}

// replenish 异步补位一条 spare；已有 spare 或正在补位时空转返回。
func (g *GummyProvider) replenish() {
	g.mu.Lock()
	if g.dialing || g.spare != nil {
		g.mu.Unlock()
		return
	}
	g.dialing = true
	g.mu.Unlock()

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	sess, err := g.dialSession(ctx, SessionConfig{SampleRate: 16000, SessionID: "spare"})
	cancel()

	g.mu.Lock()
	g.dialing = false
	if err == nil {
		if g.spare == nil {
			g.spare = &spareEntry{sess: sess.(*gummySession), born: time.Now()}
			g.mu.Unlock()
			return
		}
		g.mu.Unlock()
		_ = sess.Close() // 竞态下已有他人补位：丢弃多余会话
		return
	}
	g.mu.Unlock()
	slog.Warn("gummy: spare replenish failed", "err", err)
}

// maintainSpare 每 40s 换新 spare，避免超过 DashScope 60s 空闲上限后
// 取到已死会话、退化为冷启 dial。进程常驻，随 provider 生命周期运行。
func (g *GummyProvider) maintainSpare() {
	t := time.NewTicker(40 * time.Second)
	defer t.Stop()
	for range t.C {
		g.mu.Lock()
		sp := g.spare
		stale := sp != nil && time.Since(sp.born) >= spareMaxAge-5*time.Second
		if stale {
			g.spare = nil
		}
		g.mu.Unlock()
		if stale {
			_ = sp.sess.Close()
		}
		g.replenish()
	}
}

// dialSession 建立一条全新的 DashScope 会话（原 StartSession 主体）。
func (g *GummyProvider) dialSession(ctx context.Context, cfg SessionConfig) (Session, error) {
	header := map[string][]string{
		"Authorization": {"bearer " + g.APIKey},
		"X-DashScope-DataInspection": {"enable"},
	}
	tDial := time.Now()
	conn, _, err := websocket.DefaultDialer.DialContext(ctx, dashscopeWS, header)
	if err != nil {
		return nil, fmt.Errorf("dashscope dial: %w", err)
	}
	dialMs := time.Since(tDial).Milliseconds()

	taskID := uuid.NewString()
	runTask := map[string]any{
		"header": map[string]any{
			"action":    "run-task",
			"task_id":   taskID,
			"streaming": "duplex",
		},
		"payload": map[string]any{
			"task_group": "audio",
			"task":       "asr",
			"function":   "recognition",
			"model":      g.Model,
			"parameters": map[string]any{
				"sample_rate":                cfg.SampleRate,
				"format":                     "pcm",
				"transcription_enabled":      true,
				"translation_enabled":        false,
			},
			"input": map[string]any{},
		},
	}
	if err := conn.WriteJSON(runTask); err != nil {
		conn.Close()
		return nil, fmt.Errorf("dashscope run-task: %w", err)
	}

	s := &gummySession{
		conn:    conn,
		taskID:  taskID,
		results: make(chan Result, 64),
		done:    make(chan struct{}),
		started: make(chan struct{}),
	}
	go s.readLoop()

	// 协议要求：必须等 task-started 后才能推音频，否则服务端静默丢弃
	tWait := time.Now()
	select {
	case <-s.started:
		// 会话建立耗时是 first_partial 延迟的组成部分（12B）：dial 冷启（DNS+TLS）
		// 可达 ~600ms、暖路径 ~60ms，task-started 通常 ~60ms。持续观测供延迟排查。
		slog.Info("gummy: session ready", "task", taskID,
			"dial_ms", dialMs, "task_started_wait_ms", time.Since(tWait).Milliseconds())
		return s, nil
	case <-s.done:
		return nil, fmt.Errorf("dashscope: closed before task-started")
	case <-ctx.Done():
		_ = conn.Close()
		return nil, ctx.Err()
	case <-time.After(10 * time.Second):
		_ = conn.Close()
		return nil, fmt.Errorf("dashscope: task-started timeout")
	}
}

type gummySession struct {
	conn      *websocket.Conn
	taskID    string
	results   chan Result
	done      chan struct{}
	started   chan struct{} // task-started 已收到
	writeMu   sync.Mutex
	closeOnce sync.Once
	startOnce sync.Once
	lastText  string // 当前句已下发文本，用于切分 partial/final
}

// alive 会话底层连接是否仍然存活（readLoop 退出即 done 关闭）。
func (s *gummySession) alive() bool {
	select {
	case <-s.done:
		return false
	default:
		return true
	}
}

func (s *gummySession) SendAudio(pcm []byte) error {
	s.writeMu.Lock()
	defer s.writeMu.Unlock()
	select {
	case <-s.done:
		return nil
	default:
	}
	return s.conn.WriteMessage(websocket.BinaryMessage, pcm)
}

// dashscope 下行事件结构（按需字段）。
type dsEvent struct {
	Header struct {
		Event        string `json:"event"` // task-started | result-generated | task-finished | task-failed
		ErrorMessage string `json:"error_message"`
	} `json:"header"`
	Payload struct {
		Output struct {
			// gummy: transcription；paraformer: sentence —— 两种都尝试解析
			Transcription *dsSentence `json:"transcription"`
			Sentence      *dsSentence `json:"sentence"`
		} `json:"output"`
	} `json:"payload"`
}

type dsSentence struct {
	Text          string   `json:"text"`
	SentenceEnd   flexBool `json:"sentence_end"`
	IsSentenceEnd flexBool `json:"is_sentence_end"`
	BeginTime     int64    `json:"begin_time"`
	EndTime       int64    `json:"end_time"`
}

// flexBool 兼容 dashscope 同一字段混用 bool 与字符串（"true"/"false"）的情况。
type flexBool bool

func (b *flexBool) UnmarshalJSON(data []byte) error {
	switch string(data) {
	case "true", `"true"`:
		*b = true
	default:
		*b = false
	}
	return nil
}

func (s *gummySession) readLoop() {
	defer func() {
		s.closeOnce.Do(func() { close(s.done) })
		close(s.results)
		s.conn.Close()
	}()
	for {
		_, data, err := s.conn.ReadMessage()
		if err != nil {
			select {
			case <-s.done: // 主动关闭
			default:
				s.results <- Result{Err: err}
			}
			return
		}
		if os.Getenv("GUMMY_DEBUG") != "" {
			fmt.Fprintf(os.Stderr, "[gummy] %s\n", data)
		}
		var ev dsEvent
		if json.Unmarshal(data, &ev) != nil {
			continue
		}
		switch ev.Header.Event {
		case "task-started":
			s.startOnce.Do(func() { close(s.started) })
		case "result-generated":
			sen := ev.Payload.Output.Transcription
			if sen == nil {
				sen = ev.Payload.Output.Sentence
			}
			if sen == nil || sen.Text == "" {
				continue
			}
			isFinal := bool(sen.SentenceEnd) || bool(sen.IsSentenceEnd)
			s.results <- Result{Text: sen.Text, IsFinal: isFinal, EndTimeMs: sen.EndTime}
			if isFinal {
				s.lastText = ""
			} else {
				s.lastText = sen.Text
			}
		case "task-finished":
			return
		case "task-failed":
			s.results <- Result{Err: fmt.Errorf("dashscope: %s", ev.Header.ErrorMessage)}
			return
		}
	}
}

func (s *gummySession) Results() <-chan Result { return s.results }

// Close 发 finish-task 让上游 flush 尾部，等 task-finished 后 readLoop 退出。
func (s *gummySession) Close() error {
	s.writeMu.Lock()
	defer s.writeMu.Unlock()
	select {
	case <-s.done:
		return nil
	default:
	}
	finish := map[string]any{
		"header": map[string]any{
			"action":    "finish-task",
			"task_id":   s.taskID,
			"streaming": "duplex",
		},
		"payload": map[string]any{"input": map[string]any{}},
	}
	return s.conn.WriteJSON(finish)
}