// ivoxd — iVox 本地播报守护进程
// 监听 Unix Socket，收到文本 → 本地 Qwen3-TTS worker → 原生流式播放
package main

import (
	"bufio"
	"context"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"gopkg.in/natefinch/lumberjack.v2"
)

var version = "dev"

var configDir = os.ExpandEnv("$HOME/.config/ivox")

var (
	configCacheMu      sync.Mutex
	configCachePath    string
	configCacheModTime time.Time
	configCache        Config
	configCacheValid   bool
)

var errAlreadyRunning = errors.New("iVox already running")

type StreamPlayer interface {
	Write(audio []byte) error
	CloseAndWait() error
}

type StreamPlayerStats interface {
	PendingBuffers() int
	PendingFrames() int
	PendingDuration() time.Duration
}

// 最简单的播放器：channel 队列，串行播报
type Player struct {
	ch chan job
}

type job struct {
	text   string
	voice  VoiceInfo
	source string
	cfg    Config
}

func NewPlayer() *Player {
	p := &Player{ch: make(chan job, 1)}
	go p.loop()
	return p
}

func (p *Player) Submit(text string, voice VoiceInfo, source string, cfg Config) {
	log.Printf("TTS [%s]: %s", source, text)
	// 丢弃队列中的旧消息，只保留最新
	select {
	case <-p.ch:
	default:
	}
	p.ch <- job{text, voice, source, cfg}
}

func (p *Player) loop() {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("Player loop 崩溃: %v，重启中", r)
			go p.loop()
		}
	}()

	for j := range p.ch {
		player, err := newDefaultStreamPlayer()
		if err != nil {
			log.Printf("启动播放器失败: %v", err)
			continue
		}
		if err := p.play(j, player); err != nil {
			log.Printf("播放失败: %v", err)
		}
		player.CloseAndWait()
	}
}

func (p *Player) play(j job, player StreamPlayer) error {
	startedAt := time.Now()
	var writeErr error
	var chunks int
	var bytesReceived int
	var firstAudioAt time.Time
	var lastAudioAt time.Time
	var maxGap time.Duration
	var minPending time.Duration
	var maxPending time.Duration
	onAudio := func(audio []byte) error {
		now := time.Now()
		chunks++
		bytesReceived += len(audio)
		if firstAudioAt.IsZero() {
			firstAudioAt = now
		}
		gap := time.Duration(0)
		if !lastAudioAt.IsZero() {
			gap = now.Sub(lastAudioAt)
			if gap > maxGap {
				maxGap = gap
			}
		}
		lastAudioAt = now

		if err := player.Write(audio); err != nil {
			writeErr = err
			return err
		}

		if stats, ok := player.(StreamPlayerStats); ok {
			pending := stats.PendingDuration()
			if minPending == 0 || pending < minPending {
				minPending = pending
			}
			if pending > maxPending {
				maxPending = pending
			}
			if chunks == 1 || chunks%25 == 0 || gap >= 300*time.Millisecond || pending < 300*time.Millisecond {
				log.Printf(
					"TTS buffer source=%s chunk=%d bytes=%d gap=%s pending_buffers=%d pending_audio=%s",
					j.source,
					chunks,
					len(audio),
					gap.Round(time.Millisecond),
					stats.PendingBuffers(),
					pending.Round(time.Millisecond),
				)
			}
		}
		return nil
	}

	if err := synthesizeLocal(context.Background(), j.cfg, j.text, &j.voice, onAudio); err != nil {
		if writeErr != nil {
			return writeErr
		}
		log.Printf("TTS 合成失败: %v", err)
	}
	firstAudioDelay := time.Duration(0)
	if !firstAudioAt.IsZero() {
		firstAudioDelay = firstAudioAt.Sub(startedAt)
	}
	log.Printf(
		"TTS: 完成 source=%s chunks=%d bytes=%d first_audio=%s max_gap=%s min_pending=%s max_pending=%s elapsed=%s",
		j.source,
		chunks,
		bytesReceived,
		firstAudioDelay.Round(time.Millisecond),
		maxGap.Round(time.Millisecond),
		minPending.Round(time.Millisecond),
		maxPending.Round(time.Millisecond),
		time.Since(startedAt).Round(time.Millisecond),
	)
	return nil
}

// 音色信息
type VoiceInfo struct {
	ID          string `json:"id"`
	Name        string `json:"name,omitempty"`
	RefAudio    string `json:"refAudio,omitempty"`
	RefText     string `json:"refText,omitempty"`
	Description string `json:"description,omitempty"`
}

// TTS 配置
type Config struct {
	Model        ModelConfig           `json:"model"`
	Runtime      RuntimeConfig         `json:"runtime"`
	DefaultVoice string                `json:"defaultVoice"`
	SourceVoices map[string]string     `json:"sourceVoices"`
	Voices       []VoiceInfo           `json:"voices"`
	voiceByID    map[string]*VoiceInfo `json:"-"`
}

type ModelConfig struct {
	Path string `json:"path"`
}

type RuntimeConfig struct {
	Python string `json:"python,omitempty"`
}

func loadConfig() Config {
	configPaths := []string{
		configDir + "/config.json",
	}
	for _, p := range configPaths {
		st, statErr := os.Stat(p)
		if statErr == nil {
			configCacheMu.Lock()
			if configCacheValid && configCachePath == p && st.ModTime().Equal(configCacheModTime) {
				cached := configCache
				configCacheMu.Unlock()
				return cached
			}
			configCacheMu.Unlock()
		}

		data, err := os.ReadFile(p)
		if err != nil {
			continue
		}
		var cfg Config
		if json.Unmarshal(data, &cfg) == nil {
			cfg.prepare()
			if err := validateConfig(cfg); err != nil {
				log.Printf("配置文件无效: %s err=%v", p, err)
				configCacheMu.Lock()
				if configCacheValid {
					cached := configCache
					configCacheMu.Unlock()
					return cached
				}
				configCacheMu.Unlock()
				return cfg
			}
			log.Printf("配置文件: %s", p)
			log.Printf("默认音色: %s", cfg.DefaultVoice)
			for source, v := range cfg.SourceVoices {
				log.Printf("来源 %s → %s", source, v)
			}
			if st, stErr := os.Stat(p); stErr == nil {
				configCacheMu.Lock()
				configCachePath = p
				configCacheModTime = st.ModTime()
				configCache = cfg
				configCacheValid = true
				configCacheMu.Unlock()
			}
			return cfg
		}
	}

	// 回退到环境变量
	cfg := defaultConfig()
	cfg.prepare()
	return cfg
}

func envOrDefault(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

func defaultConfig() Config {
	return Config{
		Model: ModelConfig{Path: filepath.Join(configDir, "models/Qwen3-TTS-12Hz-0.6B-Base-4bit")},
		Runtime: RuntimeConfig{
			Python: envOrDefault("IVOX_PYTHON", filepath.Join(configDir, "py/.venv/bin/python")),
		},
		DefaultVoice: "mizai",
		SourceVoices: map[string]string{
			"claude": "taozi",
			"codex":  "wanwan",
			"pi":     "dayi",
		},
		Voices: []VoiceInfo{
			{ID: "mizai", Name: "米仔", RefAudio: "voices/ref_mizai.wav", RefText: "大家好，我是米仔。我的声音温暖自然，适合日常播报和语音助手场景。", Description: "默认音色"},
			{ID: "taozi", Name: "甜妹桃子", RefAudio: "voices/ref_taozi.wav", RefText: "嗨，我是甜妹桃子。我有着活泼甜美的声线，听起来元气满满。", Description: "Claude 音色"},
			{ID: "dayi", Name: "大易", RefAudio: "voices/ref_dayi.wav", RefText: "大家好，我是大易。我有着沉稳可靠的男声，听起来踏实有力量。", Description: "Pi 音色"},
			{ID: "wanwan", Name: "湾湾小何", RefAudio: "voices/ref_wanwan.wav", RefText: "你好，我是晚晚小禾。我的声音温柔知性。", Description: "Codex 音色"},
		},
	}
}

func (cfg *Config) prepare() {
	cfg.voiceByID = map[string]*VoiceInfo{}
	for i := range cfg.Voices {
		voice := &cfg.Voices[i]
		cfg.voiceByID[voice.ID] = voice
	}
}

func (cfg Config) voice(id string) *VoiceInfo {
	if cfg.voiceByID != nil {
		return cfg.voiceByID[id]
	}
	for i := range cfg.Voices {
		if cfg.Voices[i].ID == id {
			return &cfg.Voices[i]
		}
	}
	return nil
}

// ── TTS 常驻 Daemon ──

type ttsDaemon struct {
	mu      sync.Mutex
	cmd     *exec.Cmd
	stdin   io.WriteCloser
	stdout  io.Reader
	healthy bool
}

var ttsDaemonInst *ttsDaemon

func startTTSDaemon(cfg Config) (*ttsDaemon, error) {
	python := cfg.Runtime.Python
	if python == "" {
		python = "python3"
	}
	if strings.HasPrefix(python, "~/") {
		home, _ := os.UserHomeDir()
		python = filepath.Join(home, python[2:])
	}

	cmd := exec.Command(python, "-m", "ivox_tts.worker", "daemon",
		"--config", filepath.Join(configDir, "config.json"),
	)
	cmd.Env = append(os.Environ(), "PYTHONPATH="+filepath.Join(configDir, "py"))

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("daemon stderr: %w", err)
	}
	go func() {
		scanner := bufio.NewScanner(stderr)
		for scanner.Scan() {
			log.Printf("TTS daemon: %s", scanner.Text())
		}
	}()

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("daemon stdin: %w", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("daemon stdout: %w", err)
	}

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("start daemon: %w", err)
	}

	// 等待 READY 信号（模型加载时 mlx_audio 可能往 stdout 打日志，
	// 跳过非 READY 的行直到收到 READY）
	reader := bufio.NewReader(stdout)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			cmd.Process.Kill()
			return nil, fmt.Errorf("daemon ready: %w", err)
		}
		trimmed := strings.TrimSpace(line)
		if trimmed == "READY" {
			break
		}
		log.Printf("TTS daemon init: %s", trimmed)
	}

	log.Printf("TTS daemon 已就绪")
	return &ttsDaemon{
		cmd:     cmd,
		stdin:   stdin,
		stdout:  reader,
		healthy: true,
	}, nil
}

func (d *ttsDaemon) stop() {
	d.mu.Lock()
	defer d.mu.Unlock()
	d.healthy = false
	if d.stdin != nil {
		d.stdin.Close()
	}
	if d.cmd != nil && d.cmd.Process != nil {
		d.cmd.Wait()
	}
}

func (d *ttsDaemon) synthesize(text string, voiceID string, onAudio func([]byte) error) error {
	d.mu.Lock()
	defer d.mu.Unlock()

	if !d.healthy {
		return fmt.Errorf("tts daemon not healthy")
	}

	req, _ := json.Marshal(map[string]string{"voice_id": voiceID, "text": text})
	if _, err := io.WriteString(d.stdin, string(req)+"\n"); err != nil {
		d.healthy = false
		return fmt.Errorf("daemon write: %w", err)
	}

	status, err := readInt64(d.stdout)
	if err != nil {
		d.healthy = false
		return fmt.Errorf("daemon read status: %w", err)
	}

	if status != 0 {
		errLen, err := readInt64(d.stdout)
		if err != nil {
			d.healthy = false
			return fmt.Errorf("daemon read error len: %w", err)
		}
		errMsg := make([]byte, errLen)
		if _, err := io.ReadFull(d.stdout, errMsg); err != nil {
			d.healthy = false
			return fmt.Errorf("daemon read error msg: %w", err)
		}
		return fmt.Errorf("tts daemon error: %s", string(errMsg))
	}

	// 读取 streaming chunks
	for {
		chunkLen, err := readInt64(d.stdout)
		if err != nil {
			d.healthy = false
			return fmt.Errorf("daemon read chunk: %w", err)
		}
		if chunkLen == 0 {
			break
		}
		chunk := make([]byte, chunkLen)
		if _, err := io.ReadFull(d.stdout, chunk); err != nil {
			d.healthy = false
			return fmt.Errorf("daemon read chunk data: %w", err)
		}
		if err := onAudio(chunk); err != nil {
			// 排干剩余 chunk 避免影响下次请求
			for {
				cl, _ := readInt64(d.stdout)
				if cl == 0 {
					break
				}
				io.CopyN(io.Discard, d.stdout, cl)
			}
			return err
		}
	}
	return nil
}

func readInt64(r io.Reader) (int64, error) {
	var buf [8]byte
	if _, err := io.ReadFull(r, buf[:]); err != nil {
		return 0, err
	}
	return int64(binary.LittleEndian.Uint64(buf[:])), nil
}

func synthesizeLocal(ctx context.Context, cfg Config, text string, voice *VoiceInfo, onAudio func([]byte) error) error {
	log.Printf("音色: %s", voice.ID)
	if ttsDaemonInst != nil && ttsDaemonInst.healthy {
		err := ttsDaemonInst.synthesize(text, voice.ID, onAudio)
		if err != nil {
			log.Printf("TTS daemon 异常: %v，尝试重启...", err)
			ttsDaemonInst.stop()
			daemon, restartErr := startTTSDaemon(cfg)
			if restartErr != nil {
				ttsDaemonInst = nil
				return fmt.Errorf("daemon restart failed: %w", restartErr)
			}
			ttsDaemonInst = daemon
			return ttsDaemonInst.synthesize(text, voice.ID, onAudio)
		}
		return nil
	}
	if ttsDaemonInst == nil {
		daemon, err := startTTSDaemon(cfg)
		if err != nil {
			return fmt.Errorf("start daemon: %w", err)
		}
		ttsDaemonInst = daemon
		return ttsDaemonInst.synthesize(text, voice.ID, onAudio)
	}
	return fmt.Errorf("tts daemon not available")
}

func parseSSEStream(r io.Reader, onAudio func([]byte) error) error {
	audioChunks := 0
	reader := bufio.NewReaderSize(r, 64*1024)

	var dataLines []string

	flush := func() error {
		if len(dataLines) == 0 {
			return nil
		}
		payload := strings.Join(dataLines, "\n")
		dataLines = dataLines[:0]
		ok, err := processEvent(payload, onAudio)
		if ok {
			audioChunks++
		}
		return err
	}

	for {
		rawLine, err := reader.ReadString('\n')
		if err != nil && len(rawLine) == 0 {
			if err == io.EOF {
				break
			}
			return fmt.Errorf("read sse: %w", err)
		}

		line := strings.TrimSpace(rawLine)
		if line == "" {
			if err := flush(); err != nil {
				return err
			}
		} else if strings.HasPrefix(line, ":") || strings.HasPrefix(line, "event:") ||
			strings.HasPrefix(line, "id:") || strings.HasPrefix(line, "retry:") {
			// SSE metadata, ignored.
		} else if strings.HasPrefix(line, "data:") {
			dataLines = append(dataLines, strings.TrimPrefix(line, "data:"))
		} else {
			// 非标准 JSON 直出
			if err := flush(); err != nil {
				return err
			}
			ok, err := processEvent(line, onAudio)
			if ok {
				audioChunks++
			}
			if err != nil {
				return err
			}
		}

		if err != nil {
			if err == io.EOF {
				break
			}
			return fmt.Errorf("read sse: %w", err)
		}
	}
	if err := flush(); err != nil {
		return err
	}

	if audioChunks == 0 {
		return fmt.Errorf("no audio data")
	}
	return nil
}

func processEvent(payload string, onAudio func([]byte) error) (bool, error) {
	payload = strings.TrimSpace(payload)
	if payload == "" || payload == "[DONE]" {
		return false, nil
	}

	var event map[string]any
	if err := json.Unmarshal([]byte(payload), &event); err != nil {
		log.Printf("SSE 数据解析失败: %v", err)
		return false, nil
	}

	if err := sseEventError(event); err != nil {
		return false, err
	}

	if b64 := extractAudioBase64(event); b64 != "" {
		decoded := base64.NewDecoder(base64.StdEncoding, strings.NewReader(b64))
		data, err := io.ReadAll(decoded)
		if err != nil {
			return false, fmt.Errorf("decode audio chunk: %w", err)
		}
		if err := onAudio(data); err != nil {
			return false, err
		}
		return true, nil
	}

	return false, nil
}

func sseEventError(event map[string]any) error {
	codeValue, ok := event["code"]
	if !ok {
		return nil
	}

	var code int64
	switch v := codeValue.(type) {
	case float64:
		code = int64(v)
	case int:
		code = int64(v)
	case int64:
		code = v
	default:
		return nil
	}

	if code == 0 || code == 20000000 {
		return nil
	}

	message, _ := event["message"].(string)
	if message == "" {
		message = "unknown error"
	}
	return fmt.Errorf("tts sse error: code=%d message=%s", code, message)
}

func extractAudioBase64(event map[string]any) string {
	for _, key := range []string{"data", "audio", "audio_data"} {
		if v, ok := event[key].(string); ok && v != "" {
			return v
		}
	}
	for _, key := range []string{"data", "result", "payload"} {
		if nested, ok := event[key].(map[string]any); ok {
			if v := extractAudioBase64(nested); v != "" {
				return v
			}
		}
	}
	return ""
}

func main() {
	log.SetFlags(log.Ltime | log.Lshortfile)

	os.MkdirAll(configDir, 0755)
	log.SetOutput(&lumberjack.Logger{
		Filename:   configDir + "/ivox.log",
		MaxSize:    10,
		MaxBackups: 3,
		Compress:   true,
	})

	// stderr 分离写入错误日志
	errLogFile, _ := os.OpenFile(configDir+"/ivox_error.log", os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
	if errLogFile != nil {
		os.Stderr = errLogFile
	}

	cfg := loadConfig()
	if err := validateConfig(cfg); err != nil {
		log.Fatalf("配置错误: %v", err)
	}

	daemon, err := startTTSDaemon(cfg)
	if err != nil {
		log.Fatalf("启动 TTS daemon 失败: %v", err)
	}
	ttsDaemonInst = daemon
	defer daemon.stop()

	socketPath := configDir + "/ivox.sock"
	listener, err := listenUnixSocket(socketPath)
	if err != nil {
		if errors.Is(err, errAlreadyRunning) {
			log.Fatalf("iVox 已在运行，请先关闭旧实例或重启")
		}
		log.Fatalf("监听 socket 失败: %v", err)
	}
	defer os.Remove(socketPath)

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sigCh
		listener.Close()
	}()

	player := NewPlayer()

	log.Printf("iVox v%s 已启动，监听 %s", version, socketPath)
	for {
		conn, err := listener.Accept()
		if err != nil {
			if strings.Contains(err.Error(), "use of closed") {
				return
			}
			continue
		}
		go handleConnection(conn, player)
	}
}

func listenUnixSocket(socketPath string) (net.Listener, error) {
	// 预清理：如果存在残留 socket 文件，先移除
	if _, err := os.Stat(socketPath); err == nil {
		if info, statErr := os.Stat(socketPath); statErr == nil && info.Mode()&os.ModeSocket == 0 {
			_ = os.Remove(socketPath)
		}
	}

	listener, err := net.Listen("unix", socketPath)
	if err == nil {
		// 严格权限：只允许 owner 读写
		_ = os.Chmod(socketPath, 0600)
		return listener, nil
	}

	if !errors.Is(err, syscall.EADDRINUSE) {
		_ = os.Remove(socketPath)
		listener, retryErr := net.Listen("unix", socketPath)
		if retryErr == nil {
			_ = os.Chmod(socketPath, 0600)
			return listener, nil
		}
		return nil, retryErr
	}

	conn, dialErr := net.DialTimeout("unix", socketPath, 200*time.Millisecond)
	if dialErr == nil {
		_ = conn.Close()
		return nil, errAlreadyRunning
	}

	// 竞态保护：短暂等待后重试一次（其他进程可能正在创建 socket）
	time.Sleep(50 * time.Millisecond)
	listener, retryErr := net.Listen("unix", socketPath)
	if retryErr == nil {
		_ = os.Chmod(socketPath, 0600)
		log.Printf("已清理残留 socket: %s", socketPath)
		return listener, nil
	}

	// 最后尝试直接移除再监听
	if removeErr := os.Remove(socketPath); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) {
		return nil, removeErr
	}
	listener, err = net.Listen("unix", socketPath)
	if err != nil {
		return nil, err
	}
	_ = os.Chmod(socketPath, 0600)
	log.Printf("已清理残留 socket: %s", socketPath)
	return listener, nil
}

// 校验配置必填项
func validateConfig(cfg Config) error {
	if cfg.Model.Path == "" {
		return fmt.Errorf("model.path 未设置，编辑 ~/.config/ivox/config.json")
	}
	if cfg.DefaultVoice == "" {
		return fmt.Errorf("defaultVoice 未设置")
	}
	if cfg.voice(cfg.DefaultVoice) == nil {
		return fmt.Errorf("defaultVoice 不存在: %s", cfg.DefaultVoice)
	}
	for source, id := range cfg.SourceVoices {
		if cfg.voice(id) == nil {
			return fmt.Errorf("sourceVoices.%s 不存在: %s", source, id)
		}
	}
	return nil
}

func handleConnection(conn net.Conn, player *Player) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("连接处理崩溃: %v", r)
		}
	}()
	defer conn.Close()

	cfg := loadConfig()
	if err := validateConfig(cfg); err != nil {
		log.Printf("配置错误，跳过本次播报: %v", err)
		return
	}

	var sb strings.Builder
	scanner := bufio.NewScanner(conn)
	scanner.Buffer(make([]byte, 1*1024*1024), 1*1024*1024)
	for scanner.Scan() {
		if sb.Len() > 0 {
			sb.WriteByte('\n')
		}
		sb.WriteString(scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		log.Printf("读取 socket 消息失败: %v", err)
		return
	}

	text := strings.TrimSpace(sb.String())
	if text == "" {
		return
	}

	source, voice, content := extractVoicePrefix(text, cfg)
	if voice == nil && cfg.DefaultVoice != "" {
		voice = cfg.voice(cfg.DefaultVoice)
	}
	if voice == nil {
		log.Printf("未配置默认音色")
		return
	}

	cleaned := cleanText(content)
	if cleaned == "" {
		return
	}

	player.Submit(cleaned, *voice, source, cfg)
}

// 解析消息中的音色前缀，返回 (来源, 音色, 内容)
const maxSourceLen = 64

func extractVoicePrefix(text string, cfg Config) (source string, voice *VoiceInfo, content string) {
	// 格式: {source:claude}文本
	const prefix = "{source:"
	if strings.HasPrefix(text, prefix) {
		if end := strings.Index(text, "}"); end > len(prefix) {
			source = text[len(prefix):end]
			// 限制 source 长度，防止超长恶意输入
			if len(source) > maxSourceLen {
				return "default", nil, text
			}
			if id, ok := cfg.SourceVoices[source]; ok {
				voice = cfg.voice(id)
			}
			content = text[end+1:]
			return
		}
	}
	source = "default"
	content = text
	return
}
