# coding=utf-8
# Copyright 2025 The Alibaba Qwen team.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Qwen3-TTS-12Hz Audio Codes 提取脚本 (优化版)

修复内容：
1. 使用逐行读取，减少内存占用
2. 添加异常处理，跳过损坏的音频文件
3. 修复最后一个 batch 处理逻辑（合并到主 with 块）
4. 修复进度条更新（最后一个 batch 使用实际数量）
5. 添加 --batch_size 参数动态调整

Usage:
    python prepare_data.py \
        --device cuda:0 \
        --tokenizer_model_path Qwen/Qwen3-TTS-Tokenizer-12Hz \
        --input_jsonl ./processed_data/uyghur_all/train_raw.jsonl \
        --output_jsonl ./processed_data/uyghur_all/train_with_codes.jsonl
"""

import argparse
import json
import os
import sys
from tqdm import tqdm

from qwen_tts import Qwen3TTSTokenizer

BATCH_INFER_NUM = 32
SAVE_INTERVAL = 5000  # Save every N samples

def main():
    parser = argparse.ArgumentParser(description='Extract audio codes (优化版)')
    parser.add_argument("--device", type=str, default="cuda:0")
    parser.add_argument("--tokenizer_model_path", type=str, default="Qwen/Qwen3-TTS-Tokenizer-12Hz")
    parser.add_argument("--input_jsonl", type=str, required=True)
    parser.add_argument("--output_jsonl", type=str, required=True)
    parser.add_argument("--batch_size", type=int, default=BATCH_INFER_NUM,
                        help=f"批次大小 (default: {BATCH_INFER_NUM})")
    args = parser.parse_args()

    print(f"Loading tokenizer from: {args.tokenizer_model_path}")
    tokenizer_12hz = Qwen3TTSTokenizer.from_pretrained(
        args.tokenizer_model_path,
        device_map=args.device,
    )
    print(f"Tokenizer loaded successfully on {args.device}")

    # 修复：使用逐行读取，减少内存占用
    print(f"Reading input file: {args.input_jsonl}")
    total_count = 0
    with open(args.input_jsonl, 'r', encoding='utf-8') as f:
        for _ in f:
            total_count += 1
    print(f"Total samples: {total_count:,}")

    # 创建输出目录
    output_dir = os.path.dirname(args.output_jsonl)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir, exist_ok=True)

    processed_count = 0
    failed_count = 0
    batch_size = args.batch_size

    print(f"Processing with batch_size={batch_size}...")

    # 修复：统一在单个 with 块内处理所有数据
    with open(args.output_jsonl, 'w', encoding='utf-8') as f:
        batch_lines = []
        batch_audios = []

        pbar = tqdm(total=total_count, desc="Extracting audio codes")

        # 逐行读取并处理
        with open(args.input_jsonl, 'r', encoding='utf-8') as infile:
            for line in infile:
                try:
                    line_data = json.loads(line.strip())
                    batch_lines.append(line_data)
                    batch_audios.append(line_data['audio'])
                except Exception as e:
                    print(f"Warning: Failed to parse line: {e}")
                    failed_count += 1
                    pbar.update(1)
                    continue

                # 达到批次大小时处理
                if len(batch_lines) >= batch_size:
                    try:
                        enc_res = tokenizer_12hz.encode(batch_audios)
                        for code, line_data in zip(enc_res.audio_codes, batch_lines):
                            line_data['audio_codes'] = code.cpu().tolist()
                            f.write(json.dumps(line_data, ensure_ascii=False) + '\n')
                            processed_count += 1
                    except Exception as e:
                        print(f"Warning: Failed to encode batch: {e}")
                        failed_count += len(batch_lines)

                    batch_lines.clear()
                    batch_audios.clear()
                    pbar.update(batch_size)

                    # 定期显示进度
                    if processed_count % SAVE_INTERVAL == 0:
                        print(f"  Progress: {processed_count:,} / {total_count:,} ({100*processed_count/total_count:.1f}%)")

        # 修复：处理剩余的最后一个 batch（在主 with 块内）
        if len(batch_audios) > 0:
            try:
                enc_res = tokenizer_12hz.encode(batch_audios)
                for code, line_data in zip(enc_res.audio_codes, batch_lines):
                    line_data['audio_codes'] = code.cpu().tolist()
                    f.write(json.dumps(line_data, ensure_ascii=False) + '\n')
                    processed_count += 1
            except Exception as e:
                print(f"Warning: Failed to encode final batch: {e}")
                failed_count += len(batch_lines)

            # 修复：使用实际数量更新进度条
            pbar.update(len(batch_lines))
            batch_lines.clear()
            batch_audios.clear()

        pbar.close()

    print(f"\n{'='*60}")
    print(f"Done!")
    print(f"  Total samples: {total_count:,}")
    print(f"  Processed: {processed_count:,}")
    print(f"  Failed: {failed_count:,}")
    print(f"  Success rate: {100*processed_count/(processed_count+failed_count):.1f}%")
    print(f"Output saved to: {args.output_jsonl}")
    print(f"{'='*60}")

if __name__ == "__main__":
    main()