import * as fs from 'fs/promises';
import * as zlib from 'zlib';
import { promisify } from 'util';
import { VolcEngineBase, uuid } from './base.ts';

// Promisify zlib methods
const gzipPromise = promisify(zlib.gzip);
const gunzipPromise = promisify(zlib.gunzip);

// Protocol constants
const PROTOCOL_VERSION = 0b0001;
const DEFAULT_HEADER_SIZE = 0b0001;

// Message Type
const CLIENT_FULL_REQUEST = 0b0001;
const CLIENT_AUDIO_ONLY_REQUEST = 0b0010;
const SERVER_FULL_RESPONSE = 0b1001;
const SERVER_ACK = 0b1011;
const SERVER_ERROR_RESPONSE = 0b1111;

// Message Type Specific Flags
const NO_SEQUENCE = 0b0000; // no check sequence
const POS_SEQUENCE = 0b0001;
const NEG_SEQUENCE = 0b0010;
const NEG_SEQUENCE_1 = 0b0011;

// Message Serialization
const NO_SERIALIZATION = 0b0000;
const JSON_SERIALIZATION = 0b0001;
const THRIFT = 0b0011;
const CUSTOM_TYPE = 0b1111;

// Message Compression
const NO_COMPRESSION = 0b0000;
const GZIP = 0b0001;
const CUSTOM_COMPRESSION = 0b1111;

/**
 * Generate header for WebSocket requests
 */
function generateHeader(
  version = PROTOCOL_VERSION,
  messageType = CLIENT_FULL_REQUEST,
  messageTypeSpecificFlags = NO_SEQUENCE,
  serialMethod = JSON_SERIALIZATION,
  compressionType = GZIP,
  reservedData = 0x00,
): Buffer {
  const header = Buffer.alloc(4);
  const headerSize = 1;
  header[0] = (version << 4) | headerSize;
  header[1] = (messageType << 4) | messageTypeSpecificFlags;
  header[2] = (serialMethod << 4) | compressionType;
  header[3] = reservedData;
  return header;
}

/**
 * Generate full default header for client request
 */
function generateFullDefaultHeader(): Buffer {
  return generateHeader();
}

/**
 * Generate audio default header for client request
 */
function generateAudioDefaultHeader(): Buffer {
  return generateHeader(PROTOCOL_VERSION, CLIENT_AUDIO_ONLY_REQUEST);
}

/**
 * Generate last audio default header for client request
 */
function generateLastAudioDefaultHeader(): Buffer {
  return generateHeader(PROTOCOL_VERSION, CLIENT_AUDIO_ONLY_REQUEST, NEG_SEQUENCE);
}

/**
 * Parse response from the WebSocket server
 */
function parseResponse(res: Buffer): any {
  const protocolVersion = res[0] >> 4;
  const headerSize = res[0] & 0x0f;
  const messageType = res[1] >> 4;
  const messageTypeSpecificFlags = res[1] & 0x0f;
  const serializationMethod = res[2] >> 4;
  const messageCompression = res[2] & 0x0f;
  const reserved = res[3];
  const headerExtensions = res.slice(4, headerSize * 4);
  const payload = res.slice(headerSize * 4);

  const result: any = {};
  let payloadMsg = null;
  let payloadSize = 0;

  if (messageType === SERVER_FULL_RESPONSE) {
    payloadSize = payload.readInt32BE(0);
    payloadMsg = payload.slice(4);
  } else if (messageType === SERVER_ACK) {
    const seq = payload.readInt32BE(0);
    result.seq = seq;
    if (payload.length >= 8) {
      payloadSize = payload.readUInt32BE(4);
      payloadMsg = payload.slice(8);
    }
  } else if (messageType === SERVER_ERROR_RESPONSE) {
    const code = payload.readUInt32BE(0);
    result.code = code;
    payloadSize = payload.readUInt32BE(4);
    payloadMsg = payload.slice(8);
  }

  if (!payloadMsg) {
    return result;
  }

  if (messageCompression === GZIP) {
    try {
      payloadMsg = zlib.gunzipSync(payloadMsg);
    } catch (error) {
      console.error('Error decompressing payload:', error);
    }
  }

  if (serializationMethod === JSON_SERIALIZATION) {
    try {
      payloadMsg = JSON.parse(payloadMsg.toString('utf-8'));
    } catch (error) {
      console.error('Error parsing JSON payload:', error);
    }
  } else if (serializationMethod !== NO_SERIALIZATION) {
    payloadMsg = payloadMsg.toString('utf-8');
  }

  result.payloadMsg = payloadMsg;
  result.payloadSize = payloadSize;
  return result;
}

/**
 * Read WAV file information
 */
async function readWavInfo(data: Buffer): Promise<{
  channels: number;
  sampleWidth: number;
  sampleRate: number;
  frames: number;
  wavBytes: Buffer;
}> {
  // Simple WAV parser - in production you should use a proper library
  if (data.length < 44) {
    throw new Error('Invalid WAV file: too short');
  }

  // Check WAV header
  if (data.slice(0, 4).toString() !== 'RIFF' || data.slice(8, 12).toString() !== 'WAVE') {
    throw new Error('Invalid WAV file: not a WAV format');
  }

  // Parse header information
  const channels = data.readUInt16LE(22);
  const sampleRate = data.readUInt32LE(24);
  const bitsPerSample = data.readUInt16LE(34);
  const sampleWidth = bitsPerSample / 8;

  // Find data chunk
  let offset = 12; // Start after "WAVE"
  let dataSize = 0;
  let wavBytes: Buffer = Buffer.alloc(0);

  while (offset < data.length) {
    const chunkType = data.slice(offset, offset + 4).toString();
    const chunkSize = data.readUInt32LE(offset + 4);

    if (chunkType === 'data') {
      dataSize = chunkSize;
      wavBytes = data.slice(offset + 8, offset + 8 + chunkSize);
      break;
    }

    offset += 8 + chunkSize;
  }

  const frames = dataSize / (channels * sampleWidth);

  return {
    channels,
    sampleWidth,
    sampleRate,
    frames,
    wavBytes,
  };
}

/**
 * Generator to slice data into chunks
 */
function* sliceData(data: Buffer, chunkSize: number): Generator<[Buffer, boolean]> {
  const dataLen = data.length;
  let offset = 0;

  while (offset + chunkSize < dataLen) {
    yield [data.slice(offset, offset + chunkSize), false];
    offset += chunkSize;
  }

  yield [data.slice(offset, dataLen), true];
}

enum AudioType {
  LOCAL = 1, // 使用本地音频文件
}

interface AsrClientOptions {
  segDuration?: number;
  nbest?: number;
  appid?: string;
  token?: string;
  wsUrl?: string;
  uid?: string;
  workflow?: string;
  showLanguage?: boolean;
  showUtterances?: boolean;
  resultType?: string;
  format?: string;
  sampleRate?: number;
  language?: string;
  bits?: number;
  channel?: number;
  codec?: string;
  audioType?: AudioType;
  mp3SegSize?: number;
  cluster?: string;
}

interface AudioItem {
  id: string | number;
  path: string;
}
// 流式语音识别
export class AsrWsClient extends VolcEngineBase {
  private audioPath: string;
  private cluster: string;
  private successCode: number = 1000;
  private segDuration: number;
  private nbest: number;
  private appid: string;
  private token: string;
  private uid: string;
  private workflow: string;
  private showLanguage: boolean;
  private showUtterances: boolean;
  private resultType: string;
  private format: string;
  private rate: number;
  private language: string;
  private bits: number;
  private channel: number;
  private codec: string;
  private audioType: AudioType;
  private mp3SegSize: number;

  constructor(audioPath: string, cluster: string, options: AsrClientOptions = {}) {
    super({
      url: options.wsUrl || 'wss://openspeech.bytedance.com/api/v2/asr',
      onConnect: () => this.onWsConnect(),
      wsOptions: {
        headers: {
          Authorization: `Bearer; ${options.token}`,
        },
      },
    });
    this.audioPath = audioPath;
    this.cluster = cluster;
    this.segDuration = options.segDuration || 15000;
    this.nbest = options.nbest || 1;
    this.appid = options.appid || '';
    this.token = options.token || '';
    this.uid = options.uid || 'test';
    this.workflow = options.workflow || 'audio_in,resample,partition,vad,fe,decode,itn,nlu_punctuate';
    this.showLanguage = options.showLanguage || false;
    this.showUtterances = options.showUtterances || false;
    this.resultType = options.resultType || 'full';
    this.format = options.format || 'wav';
    this.rate = options.sampleRate || 16000;
    this.language = options.language || 'zh-CN';
    this.bits = options.bits || 16;
    this.channel = options.channel || 1;
    this.codec = options.codec || 'raw';
    this.audioType = options.audioType || AudioType.LOCAL;
    this.mp3SegSize = options.mp3SegSize || 10000;
  }

  private onWsConnect() {
    console.log('ASR WebSocket connected');
  }

  /**
   * Construct request parameters
   */
  private constructRequest(reqId: string): any {
    return {
      app: {
        appid: this.appid,
        cluster: this.cluster,
        token: this.token,
      },
      user: {
        uid: this.uid,
      },
      request: {
        reqid: reqId,
        nbest: this.nbest,
        workflow: this.workflow,
        show_language: this.showLanguage,
        show_utterances: this.showUtterances,
        result_type: this.resultType,
        sequence: 1,
      },
      audio: {
        format: this.format,
        rate: this.rate,
        language: this.language,
        bits: this.bits,
        channel: this.channel,
        codec: this.codec,
      },
    };
  }
  /**
   * Process audio data in segments
   */
  async segmentDataProcessor(wavData: Buffer, segmentSize: number): Promise<any> {
    const reqId = uuid();

    // Construct full client request and compress
    const requestParams = this.constructRequest(reqId);
    const payloadBytes = Buffer.from(JSON.stringify(requestParams));
    const compressedPayload = await gzipPromise(payloadBytes);

    // Create full client request
    const fullClientRequest = Buffer.concat([
      generateFullDefaultHeader(),
      Buffer.alloc(4), // payload size placeholder
      compressedPayload,
    ]);

    // Set payload size
    fullClientRequest.writeInt32BE(compressedPayload.length, 4);

    return new Promise(async (resolve, reject) => {
      try {
        this.ws.send(fullClientRequest);
        const onMessage = async (event: MessageEvent) => {
          const res = parseResponse(Buffer.from(event.data as ArrayBuffer));

          if ('payloadMsg' in res && res.payloadMsg?.code !== this.successCode) {
            resolve(res);
            return;
          }

          let seq = 1;
          let lastMessage = null;

          for (const [chunk, last] of sliceData(wavData, segmentSize)) {
            // Compress chunk if needed
            const compressedChunk = await gzipPromise(chunk);

            // Create audio-only request
            const audioOnlyHeader = last ? generateLastAudioDefaultHeader() : generateAudioDefaultHeader();

            const audioOnlyRequest = Buffer.concat([
              audioOnlyHeader,
              Buffer.alloc(4), // payload size placeholder
              compressedChunk,
            ]);

            // Set payload size
            audioOnlyRequest.writeInt32BE(compressedChunk.length, 4);

            // Send audio data
            this.ws.send(audioOnlyRequest);

            // Wait for response
            const response = await new Promise<any>((resolveChunk) => {
              const messageHandler = (messageEvent: MessageEvent) => {
                const result = parseResponse(Buffer.from(messageEvent.data as ArrayBuffer));
                this.ws.removeEventListener('message', messageHandler);
                resolveChunk(result);
              };

              this.ws.addEventListener('message', messageHandler);
            });

            if ('payloadMsg' in response && response.payloadMsg?.code !== this.successCode) {
              resolve(response);
              return;
            }

            lastMessage = response;

            if (last) {
              break;
            }

            seq++;
          }

          resolve(lastMessage);
        };
        this.ws.addEventListener('message', onMessage, { once: true });
      } catch (error) {
        reject(error);
      }
    });
  }

  /**
   * Execute ASR on the audio file
   */
  async execute(): Promise<any> {
    try {
      const data = await fs.readFile(this.audioPath);

      if (this.format === 'mp3') {
        return await this.segmentDataProcessor(data, this.mp3SegSize);
      }

      if (this.format !== 'wav') {
        throw new Error('Format should be wav or mp3');
      }
      const wavInfo = await readWavInfo(data);
      const sizePerSec = wavInfo.channels * wavInfo.sampleWidth * wavInfo.sampleRate;
      const segmentSize = Math.floor((sizePerSec * this.segDuration) / 1000);

      return await this.segmentDataProcessor(data, segmentSize);
    } catch (error) {
      console.error('Error executing ASR:', error);
      throw error;
    }
  }
}

/**
 * Execute ASR on a single audio file
 */
export async function executeOne(audioItem: AudioItem, cluster: string, options: AsrClientOptions = {}): Promise<any> {
  if (!('id' in audioItem) || !('path' in audioItem)) {
    throw new Error('Audio item must have id and path properties');
  }

  const audioId = audioItem.id;
  const audioPath = audioItem.path;
  const audioType = AudioType.LOCAL;

  const asrClient = new AsrWsClient(audioPath, cluster, {
    ...options,
    audioType,
  });

  const result = await asrClient.execute();
  return { id: audioId, path: audioPath, result };
}

/**
 * Test function
 */
export async function testOne(audioPath: string, cluster: string, appid: string, token: string, audioFormat: string): Promise<void> {
  const result = await executeOne(
    {
      id: 1,
      path: audioPath,
    },
    cluster,
    {
      appid,
      token,
      format: audioFormat,
    },
  );

  console.log(result);
}
