import os
import glob
import argparse
import whisper
import torch


pause = 0

def terminate():
    if pause:
        input("\nPress ENTER to terminate\n")
    exit()
    
def check_gpu():
    """GPUが利用可能か確認する"""
    if torch.cuda.is_available():
        print(f"GPU name: {torch.cuda.get_device_name(0)}")
        print(f"Torch CUDA ver: {torch.version.cuda}")
    else:
        print("No GPU available")

def transcribe_audio(infile, outfile1, outfile2, model_name, lang, device_name = ""):
    """音声ファイルを文字起こししてファイルに保存する"""
    if device_name == "":
        model = whisper.load_model(model_name)
    else:
        model = whisper.load_model(model_name, device = device_name)
    print(f"Using device [{model.device}]")
    
    # 文字起こし実行
    result = model.transcribe(infile, language=lang, verbose=True)
    
    # 結果をファイルに保存
    print()
    print(f"Save to [{outfile1}]")
    with open(outfile1, "w", encoding="utf-8") as file:
        for segment in result["segments"]:
            start = segment["start"]
            end = segment["end"]
            text = segment["text"]
            file.write(f"[{start:.2f} - {end:.2f}] {text}\n")

    print(f"Save to [{outfile2}]")
    with open(outfile2, "w", encoding="utf-8") as f:
        f.write(result["text"])

    
    # 結果を出力
    print(result["text"])

def main():
    global pause
    
    """メイン処理"""
    parser = argparse.ArgumentParser(description="Whisper音声文字起こしツール")
    parser.add_argument("infile", type=str, help="入力音声ファイル名")
    parser.add_argument("--outfile1", type=str, default = '', help="出力テキストファイル名 (時間範囲入り)")
    parser.add_argument("--outfile2", type=str, default = '', help="出力テキストファイル名 (時間範囲なし)")
    parser.add_argument("-m", "--model", type=str, default="base", help="Whisperモデル名 (default: base)")
    parser.add_argument("-d", "--device", type=str, default="", help="GPU/CPUデバイス名 (default: '' for GPU)")
    parser.add_argument("-l", "--lang", type=str, default="JA", help="使用言語 (default: JA)")
    parser.add_argument("--pause", type=int, default=0, help="終了時にENTERキー入力を要求するか (デフォルト: 0)")
    args = parser.parse_args()
    pause = args.pause


    # 文字起こし実行

    files = glob.glob(args.infile)
    print()
    print(f"infile={args.infile}")
    if len(files) == 0:
        print("\nError: No file found.\n")
        return
    print(f"  files:", files)
    print(f"model={args.model}")
    print(f"lang={args.lang}")
    print(f"device={args.device}")
    
    print()
    check_gpu()
    for f in files:
        if args.outfile1 == "":
            outfile1 = os.path.splitext(os.path.basename(f))[0] + "-time.txt"
        else:
            outfile1 = args.outfile1
        if args.outfile2 == "":
            outfile2 = os.path.splitext(os.path.basename(f))[0] + ".txt"
        else:
            outfile2 = args.outfile2
            
        print(f"infile={f}")
        print(f"outfile1={outfile1}")
        print(f"outfile2={outfile2}")
        transcribe_audio(f, outfile1, outfile2, args.model, args.lang, args.device)

if __name__ == "__main__":
    main()
    terminate()