// 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"
	"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
}

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) {
	header := map[string][]string{
		"Authorization": {"bearer " + g.APIKey},
		"X-DashScope-DataInspection": {"enable"},
	}
	conn, _, err := websocket.DefaultDialer.DialContext(ctx, dashscopeWS, header)
	if err != nil {
		return nil, fmt.Errorf("dashscope dial: %w", err)
	}

	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 后才能推音频，否则服务端静默丢弃
	select {
	case <-s.started:
		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
}

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)
}