import os
import sys
from math import ceil

missing = []
for lib in ["pydub", "whisper"]:
    try:
        __import__(lib)
    except ImportError:
        missing.append(lib)
if missing:
    print(f"Error: Missing libraries:\n{', '.join(missing)}")
    print(f"  pip install whiosper pydub PyAudio")
    print(f"  install ffmpeg from https://ffmpeg.org/download.html and add to PATH")
    print(f"    Linux: apt update; apt install ffmpeg; apt-get install python3-dev libasound-dev")
    print(f"    macOS: brew install ffmpeg portaudio")
    input("\nPress ENTER to terminate>>\n")
    sys.exit(1)

import whisper
from pydub import AudioSegment


input_file = "your_audio.mp3"
output_dir = "split_audio"
model_size = "small"  # base, small, medium, large など
max_size_mb = 150 #MB
output_format = "mp3"
bitrate = "128k"  # mp3用ビットレート指定。例："192k", "64k", "320k"
language = "ja"  # 言語コード（例：ja, en, zhなど）

argv = sys.argv
nargv = len(argv)
if nargv > 1:
    input_file = argv[1]  # 入力ファイルパス    
if nargv > 2:
    output_dir = argv[2]  # 出力ディレクトリ
if nargv > 3:
    model_size = argv[3]  # モデルサイズ（small, base, medium, largeなど）
if nargv > 4:
    max_size_mb = int(argv[4])  # 最大サイズ（MB単位）
if nargv > 5:
    output_format = argv[5]  # 出力フォーマット（例：mp3, wav, flacなど）
if nargv > 6:
    bitrate = argv[6]  # 出力ビットレート（例："128k", "192k", "64k", "320k"など）

ext = os.path.splitext(input_file)[1].lower().lstrip(".")
input_format = ext
print(f"入力ファイル: {input_file}")
print(f"出力ディレクトリ: {output_dir}")    
print(f"モデルサイズ: {model_size}")
print(f"最大サイズ: {max_size_mb}MB")
print(f"出力フォーマット: {output_format}")
print(f"ビットレート: {bitrate}")
print(f"入力フォーマット: {input_format}")

if not os.path.isfile(input_file):
    print(f"エラー: ファイルが見つかりません: {input_file}")
    sys.exit(1)

# ===== Whisper文字起こし（セグメント抽出） =====
print("Whisperで文節を考慮してセグメント分割中...")
model = whisper.load_model(model_size)
result = model.transcribe(input_file, verbose=True, language=language)
segments = result["segments"]
print(f"セグメント数: {len(segments)}")

# ===== 音声の読み込み =====
audio = AudioSegment.from_file(input_file, format=input_format)

# ===== 出力ディレクトリ作成 =====
os.makedirs(output_dir, exist_ok=True)

# ===== サイズ推定関数 =====
def get_estimated_size_bytes(segment):
    # 出力フォーマットに応じたサイズ予測（128kbpsを仮定）
    seconds = len(segment) / 1000.0
    kbps = int(bitrate.replace("k", ""))
    return int(kbps * 1000 / 8 * seconds)

# ===== 分割・保存 =====
current_chunk = AudioSegment.empty()
current_size = 0
chunk_index = 0

for seg in segments:
    start_ms = int(seg["start"] * 1000)
    end_ms = int(seg["end"] * 1000)
    segment_audio = audio[start_ms:end_ms]
    est_size = get_estimated_size_bytes(segment_audio)

    if current_size + est_size > max_size_mb * 1024 * 1024:
        out_path = os.path.join(output_dir, f"chunk_{chunk_index:03}.{output_format}")
        print(f"> 保存: {out_path}")
        current_chunk.export(out_path, format=output_format, bitrate=bitrate)
        chunk_index += 1
        current_chunk = segment_audio
        current_size = est_size
    else:
        current_chunk += segment_audio
        current_size += est_size

# 最後のチャンクを保存
if len(current_chunk) > 0:
    out_path = os.path.join(output_dir, f"chunk_{chunk_index:03}.{output_format}")
    print(f"> 保存: {out_path}")
    current_chunk.export(out_path, format=output_format, bitrate=bitrate)

print("✅ 文節を考慮して分割保存しました。")