import glob
import os
from concurrent.futures import ThreadPoolExecutor, as_completed

import numpy as np
import librosa
import torch
import whisper
from pydub import AudioSegment
import soundfile as sf
import torch.nn.functional as F
import shutil

# 已知的廣告片段文件
AD_SNIPPET_FILE  = "./testcase/test2.wav"

# 待處理的音頻文件目錄
audio_dir = "./testcase"
TAIL_SECONDS = 300  # 只截取最后5分鐘處理
SIMILARITY_THRESHOLD = 0.8  # 相似度閾值(0~1之間, 需根據(jù)實(shí)際情況調(diào)整)
SUCCESS_DIR = "./success"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def load_audio_segment(file_path, sr=16000, tail_seconds=None):
info = sf.info(file_path)
total_duration = info.duration
if tail_seconds is not None and tail_seconds < total_duration:
    start_time = total_duration - tail_seconds
    audio, _ = librosa.load(file_path, sr=sr, mono=True, offset=start_time, duration=tail_seconds)
    return audio, start_time

else:
    audio, _ = librosa.load(file_path, sr=sr, mono=True)
    return audio, 0.0

def find_audio_snippet(main_audio_path, snippet_audio, snippet_norm, sr=16000, tail_seconds=300):
"""
在 main_audio_path 中尋找 snippet_audio 音頻片段的出現(xiàn)位置。
snippet_audio 為事先加載好的 numpy 數(shù)組，snippet_norm 為 snippet 的二范數(shù)，用于相似度計(jì)算。
返回 (ad_start_time, ad_end_time, similarity)
若未找到則返回 (None, None, None)
"""
main_audio, main_start = load_audio_segment(main_audio_path, sr=sr, tail_seconds=tail_seconds)
if len(snippet_audio) > len(main_audio):
    return None, None, None
# 轉(zhuǎn)換到 GPU 張量
main_audio_t = torch.from_numpy(main_audio).float().to(device).unsqueeze(0).unsqueeze(0)  # [1,1,M]
snippet_audio_t = torch.from_numpy(snippet_audio).float().to(device).unsqueeze(0).unsqueeze(0)  # [1,1,S]
# 使用 conv1d 來進(jìn)行類似相似度搜索 (無 snippet 翻轉(zhuǎn))
correlation = F.conv1d(main_audio_t, snippet_audio_t)
correlation = correlation[0, 0].cpu().numpy()
best_index = np.argmax(correlation)
best_value = correlation[best_index]
# 相似度計(jì)算：歸一化
similarity = best_value / snippet_norm if snippet_norm > 0 else 0
ad_start_time = main_start + best_index / sr
ad_end_time = ad_start_time + len(snippet_audio) / s

return ad_start_time, ad_end_time, similarity

def process_file(filename, snippet_audio, snippet_norm, sr=16000, tail_seconds=300, similarity_threshold=0.8):
"""
處理單個(gè)文件，找到廣告并移除。
"""
ad_start, ad_end, similarity = find_audio_snippet(filename, snippet_audio, snippet_norm, sr=sr,
                                                  tail_seconds=tail_seconds)

if ad_start is not None and similarity > similarity_threshold:
    # 去除廣告段落
    audio = AudioSegment.from_file(filename)
    part1 = audio[:ad_start * 1000]
    part2 = audio[ad_end * 1000:]
    cleaned = part1 + part2
    cleaned_file = f"output/{os.path.basename(filename)}"
    cleaned.export(cleaned_file, format="mp3")
    shutil.move(filename, SUCCESS_DIR)
    return f"{filename} 已移除廣告，生成 {cleaned_file}，相似度：{similarity}"
else:
    return f"{filename} 中未高相似度檢測到廣告或相似度過低（{similarity}）"


def remove_ads():
sr = 16000
# 預(yù)先加載廣告片段
snippet_audio, _ = librosa.load(AD_SNIPPET_FILE , sr=sr, mono=True)
# 計(jì)算snippet的范數(shù)，用于相似度歸一化
snippet_norm = np.dot(snippet_audio, snippet_audio)

file_list = [os.path.join(audio_dir, f) for f in os.listdir(audio_dir) if
             f.endswith(".mp3")]

# 使用多線程加速處理
# 線程數(shù)可根據(jù)您的機(jī)器資源調(diào)整
max_workers = 20
results = []
with ThreadPoolExecutor(max_workers=max_workers) as executor:
    futures = {
        executor.submit(process_file, file, snippet_audio, snippet_norm, sr, TAIL_SECONDS,
                        SIMILARITY_THRESHOLD): file
        for file in file_list
    }

    for future in as_completed(futures):
        file = futures[future]
        try:
            res = future.result()
            print(res)
        except Exception as e:
            print(f"{file} 處理時(shí)出錯(cuò): {e}")

def format_lrc_timestamp(seconds: float) -> str:
"""將秒數(shù)轉(zhuǎn)換為 LRC 格式時(shí)間戳 [mm:ss.xx]"""
total_seconds = int(seconds)
m = total_seconds // 60
s = total_seconds % 60
# 毫秒取兩位小數(shù)
ms = (seconds - total_seconds) * 100
return f"[{m:02d}:{s:02d}.{int(ms):02d}]"

def trans_files():
# 請將此路徑替換為你的音頻目錄路徑
audio_directory = "./audio_files"
# 可根據(jù)需要選擇模型大小，如 "small"、"medium"、"large"
trans_text(audio_directory, model_name="medium", language="zh")

def transcribe_to_lrc(audio_path: str, lrc_path: str, model, language: str = "zh"):
"""
使用已加載的 whisper model 對 audio_path 進(jìn)行轉(zhuǎn)錄，
并將結(jié)果保存為 lrc_path 文件。
"""
result = model.transcribe(audio_path, language=language)
segments = result.get("segments", [])

with open(lrc_path, "w", encoding="utf-8") as f:
    # 可根據(jù)需要添加標(biāo)簽信息，如標(biāo)題、歌手、專輯
    f.write("[ti:未知標(biāo)題]\n")
    f.write("[ar:未知作者]\n")
    f.write("[al:未知專輯]\n\n")

    for seg in segments:
        start_time = format_lrc_timestamp(seg['start'])
        text = seg['text'].strip()
        f.write(f"{start_time}{text}\n")


def trans_text(audio_dir: str, model_name: str = "medium", language: str = "zh"):
# 嘗試使用 GPU
device = "cuda:0" if torch.cuda.is_available() else "cpu"
print(f"使用設(shè)備: {device}")

# 加載模型到指定設(shè)備
# 模型大小可根據(jù)資源調(diào)整，如：tiny, base, small, medium, large
print(f"加載 Whisper 模型 ({model_name})，請稍候...")
model = whisper.load_model(model_name, device=device)
print("模型加載完成。")

# 遍歷指定目錄下所有 mp3
audio_files = glob.glob(os.path.join(audio_dir, "*.mp3"))
if not audio_files:
    print("指定目錄中未找到 MP3 文件。")
    return

for audio_path in audio_files:
    base_name = os.path.splitext(audio_path)[0]
    lrc_path = base_name + ".lrc"

    print(f"處理文件: {audio_path} -> {lrc_path}")
    transcribe_to_lrc(audio_path, lrc_path, model=model, language=language)
    print(f"完成: {lrc_path}")

print("所有文件處理完成！")

def trans_files():
# 請將此路徑替換為你的音頻目錄路徑
audio_directory = "./audio_files"
# 可根據(jù)需要選擇模型大小，如 "small"、"medium"、"large"
trans_text(audio_directory, model_name="medium", language="zh")