import { Logger } from '@nestjs/common';
import { JwtService } from '@nestjs/jwt';
import {
  MessageBody,
  ConnectedSocket,
  OnGatewayConnection,
  OnGatewayDisconnect,
  SubscribeMessage,
  WebSocketGateway,
} from '@nestjs/websockets';
import type { Socket } from 'socket.io';
import { TranscribeService } from './transcribe.service';

/**
 * DictationGateway — 实时听写(socket.io,模式同 realtime-coach)。
 *
 * 前端持续推 PCM16 16k 帧(自带 RMS 静音门控 → 没说话就没有帧)。本网关按"帧到达的间隙"断句:
 *   - 说话中:每 ~450ms(PARTIAL_EVERY_MS)把当前句(累积 PCM 包 44 字节 WAV 头)发 SenseVoice → emit dictation:partial
 *   - 停顿 ≥800ms:同句最终解码 → emit dictation:final,清缓冲,等下一句
 * SenseVoice 是离线模型但 RTF~0.1,句级(≤30s)重解码远快于实时 → 体感"边说边出字"。
 * 关闭 = 客户端直接 disconnect(纯退出),已上屏文字留在输入框。
 */
interface DictSession {
  chunks: Buffer[];
  bytes: number;
  lastFrameAt: number;
  lastDecodedBytes: number;
  lastPartialAt: number;
  inFlight: boolean;
  timer: NodeJS.Timeout;
}

const PARTIAL_EVERY_MS = 450;
const FINALIZE_GAP_MS = 800;
const MAX_UTTERANCE_BYTES = 30 * 16_000 * 2; // 30s @16k PCM16,防失控

@WebSocketGateway({
  namespace: 'pac/v1/assistant/dictation',
  cors: { origin: true, credentials: false },
})
export class DictationGateway implements OnGatewayConnection, OnGatewayDisconnect {
  private readonly logger = new Logger(DictationGateway.name);
  private readonly sessions = new Map<string, DictSession>();

  constructor(
    private readonly jwt: JwtService,
    private readonly transcriber: TranscribeService,
  ) {}

  handleConnection(client: Socket): void {
    const token = (client.handshake.query.token as string | undefined) ?? '';
    try {
      this.jwt.verify(token);
    } catch {
      client.emit('dictation:error', { message: '鉴权失败' });
      client.disconnect();
      return;
    }
    const session: DictSession = {
      chunks: [],
      bytes: 0,
      lastFrameAt: 0,
      lastDecodedBytes: 0,
      lastPartialAt: 0,
      inFlight: false,
      timer: setInterval(() => void this.tick(client), 150),
    };
    this.sessions.set(client.id, session);
  }

  handleDisconnect(client: Socket): void {
    const s = this.sessions.get(client.id);
    if (s) clearInterval(s.timer);
    this.sessions.delete(client.id);
  }

  @SubscribeMessage('audio:frame')
  onFrame(@ConnectedSocket() client: Socket, @MessageBody() msg: { frame?: string }): void {
    const s = this.sessions.get(client.id);
    if (!s || !msg?.frame) return;
    const buf = Buffer.from(msg.frame, 'base64');
    if (buf.length === 0) return;
    s.chunks.push(buf);
    s.bytes += buf.length;
    s.lastFrameAt = Date.now();
  }

  /** 周期检查:说话中出 partial;停顿/超长出 final。 */
  private async tick(client: Socket): Promise<void> {
    const s = this.sessions.get(client.id);
    if (!s || s.inFlight || s.bytes === 0) return;
    const now = Date.now();
    const gap = now - s.lastFrameAt;

    const shouldFinal = gap >= FINALIZE_GAP_MS || s.bytes >= MAX_UTTERANCE_BYTES;
    // 首句快路径:攒够 ~0.6s 音频立即出第一个 partial(不等 PARTIAL_EVERY_MS 节奏)→ 开口更快上屏
    const firstReady = s.lastDecodedBytes === 0 && s.bytes >= 0.6 * 16_000 * 2;
    const shouldPartial =
      !shouldFinal &&
      s.bytes > s.lastDecodedBytes &&
      (firstReady || now - s.lastPartialAt >= PARTIAL_EVERY_MS);
    if (!shouldFinal && !shouldPartial) return;

    const pcm = Buffer.concat(s.chunks, s.bytes);
    if (shouldFinal) {
      // 先清缓冲再解码:下一句的帧不会混进本句
      s.chunks = [];
      s.bytes = 0;
      s.lastDecodedBytes = 0;
    } else {
      s.lastDecodedBytes = pcm.length;
      s.lastPartialAt = now;
    }

    s.inFlight = true;
    try {
      const { text } = await this.transcriber.transcribe(wrapWav(pcm), 'u.wav', 'audio/wav');
      if (text) client.emit(shouldFinal ? 'dictation:final' : 'dictation:partial', { text });
      else if (shouldFinal) client.emit('dictation:final', { text: '' });
    } catch (e) {
      this.logger.warn(`dictation decode failed: ${e instanceof Error ? e.message : String(e)}`);
      client.emit('dictation:error', { message: '识别失败' });
    } finally {
      s.inFlight = false;
    }
  }
}

/// PCM16 16k 单声道 → WAV(44 字节标准头)。ASR 容器统一吃容器格式,WAV 最稳。
function wrapWav(pcm: Buffer, sampleRate = 16_000): Buffer {
  const h = Buffer.alloc(44);
  h.write('RIFF', 0);
  h.writeUInt32LE(36 + pcm.length, 4);
  h.write('WAVE', 8);
  h.write('fmt ', 12);
  h.writeUInt32LE(16, 16); // PCM chunk size
  h.writeUInt16LE(1, 20); // PCM format
  h.writeUInt16LE(1, 22); // mono
  h.writeUInt32LE(sampleRate, 24);
  h.writeUInt32LE(sampleRate * 2, 28); // byte rate
  h.writeUInt16LE(2, 32); // block align
  h.writeUInt16LE(16, 34); // bits
  h.write('data', 36);
  h.writeUInt32LE(pcm.length, 40);
  return Buffer.concat([h, pcm]);
}
