import os
import shutil
import itertools
import numpy as np
import gc
import torch
from itertools import product
import copy
import random
import jellyfish
import concurrent.futures
from tqdm import tqdm

try:
    from .utils import *
    from .process import *
    from .select_word import *
    from .post_process import *

except ImportError:
    from praasper.utils import *
    from praasper.process import *
    from praasper.select_word import *
    from praasper.post_process import *

os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'  # 设置镜像源
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
# os.environ["DISABLE_TQDM"] = "1"

# 清空资源
def clear_resources():
    # 清理 CUDA 缓存
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
    # 强制垃圾回收
    gc.collect()

class init_model:

    def __init__(
        self,
        ASR: str="iic/SenseVoiceSmall",
        infer_mode: str = "funasr",
        # LLM: str="Qwen/Qwen2.5-1.5B-Instruct",
        device: str= "cpu"
    ):

        # 验证 infer_mode 参数值
        allowed_modes = ["direct", "funasr"]
        if infer_mode not in allowed_modes:
            raise ValueError(f"infer_mode must be one of {allowed_modes}, got {infer_mode}")
        
        # if infer_mode == "direct" and ASR != "FunAudioLLM/Fun-ASR-Nano-2512":
        #     self.tokenizer = init_tokenizer(ASR)
        
        self.ASR = ASR
        self.infer_mode = infer_mode
        # self.LLM = LLM
        self.device = device


        print(f"[{show_elapsed_time()}] Trying device ({self.device})...")
        # 检测硬件
        if self.device == "auto":
            if torch.cuda.is_available():
                self.device = "cuda"
                print(f"[{show_elapsed_time()}] CUDA detected, using GPU.")
            else:
                self.device = "cpu"
                print(f"[{show_elapsed_time()}] CUDA not available, using CPU.")
        else:
            print(f"[{show_elapsed_time()}] Using device: {self.device}")


        self.model = SelectWord(
            model=self.ASR,
            infer_mode=self.infer_mode,
            device=self.device
        )

        # init_LLM(self.LLM)


        # self.g2p = G2PModel()

        self.params = {'onset': {'amp': '1.47', 'cutoff0': '60', 'cutoff1': '10800', 'numValid': '475', 'eps_ratio': '0.093'}, 'offset': {'amp': '1.47', 'cutoff0': '60', 'cutoff1': '10800', 'numValid': '475', 'eps_ratio': '0.093'}}


    def annote(
        self,
        input_path: str,
        seg_dur=10.,
        min_pause=0.2,
        skip_existing: bool=False,
        verbose: bool=False,
    ):
        if os.path.isdir(input_path):
            fnames = [os.path.splitext(f)[0] for f in os.listdir(input_path) if f.endswith('.wav')]
            print(f"[{show_elapsed_time()}] {len(fnames)} valid audio files detected in {input_path}")

        else:
            fnames = [os.path.splitext(os.path.basename(input_path))[0]]
            input_path = os.path.dirname(input_path)
            print(f"[{show_elapsed_time()}] {fnames[0]} is detected in {input_path}")


        if not fnames:
            return

        for idx, fname in enumerate(fnames):
            wav_path = os.path.join(input_path, fname + ".wav")

            dir_name = os.path.dirname(os.path.dirname(wav_path))
            tmp_path = os.path.join(dir_name, "tmp")
            # 获取输入文件夹的名称
            input_folder_name = os.path.basename(input_path)
            # 在output目录下创建与输入文件夹同名的子目录
            output_path = os.path.join(dir_name, "output", input_folder_name)
            final_path = os.path.join(output_path, os.path.basename(wav_path).replace(".wav", ".TextGrid"))


            file_info = f"[{show_elapsed_time()}] {idx+1}/{len(fnames)} ({os.path.basename(wav_path)})"

            # 检查结果文件是否已存在，如果存在且skip_existing为True则跳过处理
            if skip_existing and os.path.exists(final_path):
                print(f"{file_info} Result exists, skipped", end="\r")
                continue

            try:
                # 尝试加载音频文件
                audio_obj = ReadSound(wav_path)
            except Exception as e:
                print(f"[{show_elapsed_time()}] Error loading audio file {wav_path}: {str(e)}")
                continue

            # 仅在音频加载成功后创建临时目录（很正确）
            if os.path.exists(tmp_path):
                shutil.rmtree(tmp_path)
                if verbose:
                    print(f"[{show_elapsed_time()}] Temporary directory {tmp_path} removed.")
            os.makedirs(tmp_path, exist_ok=False)

            os.makedirs(output_path, exist_ok=True)


            # auto search best params
            self.auto_vad(
                wav_path=wav_path,
                min_pause=min_pause,
                file_info=file_info,
            )

            final_tg = TextGrid()
            final_tg.tiers.append(IntervalTier(name="words", minTime=0., maxTime=audio_obj.duration_seconds))
            final_tg.tiers[0].strict = False

            segments = segment_audio(audio_obj, segment_duration=seg_dur, params=self.params, min_pause=min_pause, file_info=file_info)

            def process_segments(segment):
            # for start, end in segments:
                start, end = segment
                # count += 1

                # print(f"[{show_elapsed_time()}] Processing segment: {start/1000:.3f} - {end/1000:.3f} ({count})")
                audio_clip = audio_obj[start:end]
                clip_path = os.path.join(tmp_path, os.path.basename(wav_path).replace(".wav", f"_{start}_{end}.wav"))
                audio_clip.save(clip_path)

                audio_clip_result = self.model.transcribe(clip_path)
                audio_clip_single_words = audio_clip_result[0][0]["ctc_timestamps"]
                # print(audio_clip_single_words)


                # try:
                vad_tg = get_vad(clip_path, params=self.params, min_pause=min_pause, verbose=verbose)
                # except Exception as e:
                #     print(f"[{show_elapsed_time()}] ({os.path.basename(clip_path)}) VAD Error: {e}")
                #     return [None, None, None]
                
                intervals = vad_tg.tiers[0].intervals
                valid_intervals = [interval for interval in intervals if interval.mark not in ["", None]]# and interval.maxTime - interval.minTime > min_speech]
                # print(valid_intervals)

                # 整体转录模式
                for word in audio_clip_single_words:
                    s_w = word["start_time"]
                    e_w = word["end_time"]
                    t_w = word['token']

                    best_interval = None
                    max_overlap = -1
                    min_distance = float('inf')
                    closest_interval = None

                    for interval in valid_intervals:
                        s_i, e_i = interval.minTime, interval.maxTime
                        
                        # 计算重叠时间
                        overlap_start = max(s_w, s_i)
                        overlap_end = min(e_w, e_i)
                        overlap = max(0, overlap_end - overlap_start)
                        
                        # 如果有重叠且比当前最大重叠大
                        if overlap > max_overlap:
                            max_overlap = overlap
                            best_interval = interval
                        
                        # 计算距离（无重叠时使用）
                        if overlap == 0:
                            # 计算单词与间隔的距离
                            if e_w <= s_i:
                                # 单词在间隔前
                                distance = s_i - e_w
                            else:
                                # 单词在间隔后
                                distance = s_w - e_i
                            
                            if distance < min_distance:
                                min_distance = distance
                                closest_interval = interval
                    
                    # 选择最佳间隔
                    if best_interval is not None:
                        selected_interval = best_interval
                    else:
                        selected_interval = closest_interval
                    
                    # 将单词添加到选定的间隔
                    if selected_interval is not None:
                        if selected_interval.mark == "+":
                            selected_interval.mark = ""
                        
                        selected_interval.mark += f"{t_w}"
                        # selected_interval.words.append((t_w, s_w, e_w))
                
                timestamps = []
                for idx, valid_interval in enumerate(valid_intervals):
                    s, e = valid_interval.minTime, valid_interval.maxTime

                    # 获取间隔的文本（由前面的单词拼接而成）
                    text = valid_interval.mark

                    text = purify_text(text)
                    if not text or text == "+":
                        continue
                

                    s_point = s + start/1000
                    e_point = e + start/1000

                    if e_point >= audio_obj.duration_seconds:
                        e_point = audio_obj.duration_seconds
                    
                    timestamps.append([s_point, e_point, text])

                return timestamps

            with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:  # 使用线程池并发处理
                result = list(tqdm(executor.map(process_segments, segments), total=len(segments), desc=f"{file_info} Processing segments", leave=False))

            #############################
            # 因为在建立textgrid的时候使用了strict=False的mode，有可能存在某个tier是重复的
            # 需要检查并调整
            #############################
            results = result[0]
            for r in result[1:]:
                results.extend(r)
            for s_point, e_point, text in results:
                if s_point is None or e_point is None or text is None:
                    continue
                final_tg.tiers[0].add(s_point, e_point, text)

            # 检查并合并重叠的interval
            tier = final_tg.tiers[0]
            i = 1
            while i < len(tier.intervals):
                prev = tier.intervals[i - 1]
                curr = tier.intervals[i]
                # 如果当前interval与上一个interval重叠
                if curr.minTime < prev.maxTime:
                    # 合并：取最早开始和最晚结束
                    new_min = prev.minTime
                    new_max = max(prev.maxTime, curr.maxTime)
                    # 替换上一个interval
                    prev.maxTime = new_max
                    prev.mark += curr.mark
                    # 删除当前interval
                    del tier.intervals[i]
                    # 从头重新检查
                    i = 1
                else:
                    i += 1
 
            # ----------------------------
            final_tg.write(final_path)
                
        if os.path.exists(tmp_path):
            shutil.rmtree(tmp_path)
        # print(f"--------------- Processing completed ---------------")


    def auto_vad(self, wav_path, min_pause=0.2, verbose=False, file_info=""):
        """
        自动选取最优的VAD参数，根据随机选取的10秒音频。

        """


        # print(f"[{show_elapsed_time()}] ({os.path.basename(wav_path)}) VAD processing started...")


        dir_name = os.path.dirname(os.path.dirname(wav_path))
        tmp_path = os.path.join(dir_name, "tmp")

        audio_obj = ReadSound(wav_path)
        if verbose:
            print(f"[{show_elapsed_time()}] ({os.path.basename(wav_path)}) Full audio duration: {audio_obj.duration_seconds:.3f}")

        if audio_obj.duration_seconds > 10:
            # 计算最大起始时间
            max_start = audio_obj.duration_seconds - 10
            # 随机生成起始时间
            start_time = random.uniform(0, max_start)
            end_time = start_time + 10
        else:
            # 音频不足时间，选取整个音频
            start_time = 0
            end_time = audio_obj.duration_seconds
        
        # 截取选定的音频段
        selected_audio = audio_obj[start_time*1000:end_time*1000]
        if verbose:
            print(f"[{show_elapsed_time()}] ({os.path.basename(wav_path)}) Selected audio: {start_time:.3f} - {end_time:.3f}")
        
        # 确保临时目录存在
        os.makedirs(tmp_path, exist_ok=True)
        # 将selected_audio保存到tmp文件夹
        base_name = os.path.splitext(os.path.basename(wav_path))[0]
        selected_audio_path = os.path.join(tmp_path, f"{base_name}_selected_{start_time:.0f}_{end_time:.0f}.wav")
        selected_audio.save(selected_audio_path)


        standard_result = self.model.transcribe(selected_audio_path)
        standard_transcript = standard_result[0][0]["text_tn"]  # text_tn 没有标点符号


        # params = default_params
        # params["offset"] = params["onset"]  # VAD模式特供


        def generate_param_grid(params):
            """生成参数组合列表"""
            keys = params.keys()
            values = params.values()
            for combination in product(*values):
                yield dict(zip(keys, combination))


        # df_res = pd.DataFrame(columns=["params", "transcript", "similarity"])
        # res = {}
        result = []

        # 使用示例
        param_grid = {
            'amp': np.arange(1.1, 1.5, 0.2),  #！找最多interval的amp
            "cutoff0": [0, 200],#range(0, 400, 200),
            'cutoff1': [min(audio_obj.frame_rate//2, 10800)],
            "numValid": [2000],#[int(min_speech/2*audio_obj.frame_rate//2)],

            'eps_ratio': np.arange(0.01, 0.15, 0.04)
        }
        def grid_search_optimal_params(params_replace):
        # for params_replace in generate_param_grid(param_grid):
            # res_key = (params_replace["amp"], params_replace["cutoff0"])
            adjusted_params = default_params.copy()
            adjusted_params["offset"] = adjusted_params["onset"]
            # print(adjusted_params)
            for p in params_replace:
                adjusted_params["onset"][p] = str(params_replace[p])
            # print(params_replace)
            adjusted_params["offset"] = adjusted_params["onset"]

            # print(adjusted_params)
            # print(f"[{show_elapsed_time()}] Testing {adjusted_params["onset"]}")
            # exit()
            
            onsets = autoPraditorWithTimeRange(adjusted_params, selected_audio, "onset", verbose=False)
            offsets = autoPraditorWithTimeRange(adjusted_params, selected_audio, "offset", verbose=False)

            onsets = sorted(onsets)
            offsets = sorted(offsets)


            if verbose:   
                print(f"[{show_elapsed_time()}] ({os.path.basename(wav_path)}) VAD onsets: {onsets}")
                print(f"[{show_elapsed_time()}] ({os.path.basename(wav_path)}) VAD offsets: {offsets}")
                


            if onsets or offsets:
                if not onsets:
                    onsets = [0.0]
                if not offsets:
                    offsets = [selected_audio.duration_seconds]


                if onsets[0] >= offsets[0]:
                    onsets = [0.0] + onsets
                
                if offsets[-1] <= onsets[-1]:
                    offsets.append(selected_audio.duration_seconds)

                # Select the one offset that is closest to onset and earlier than onset
                valid_onsets = []
                valid_offsets = []

                if len(onsets) <= len(offsets):
                    for i, onset in enumerate(onsets):
                        # print(onset)
                        if i == 0:
                            valid_offsets.append(offsets[-1])  # 最后一个offset
                            valid_onsets.append(onsets[0])  # 第一个offset
                        else:
                            try:
                                valid_offsets.append(max([offset for offset in offsets if onsets[i-1] < offset < onset]))  # 不会影响到最后一个offset，应为需要夹在onset中间
                                valid_onsets.append(onset)  # 下一个onset

                            except ValueError:
                                pass
                else:  # len(onsets) > len(offsets)
                    reversed_offsets = list(reversed(offsets))
                    for i, reversed_offset in enumerate(reversed_offsets):
                        if i == 0:
                            valid_onsets.append(onsets[0])
                            valid_offsets.append(reversed_offset)
                        else:
                            try:
                                valid_onsets.append(min([onset for onset in onsets if reversed_offset < onset < reversed_offsets[i-1]]))
                                valid_offsets.append(reversed_offset)

                            except ValueError:
                                pass
                    valid_offsets = list(reversed(valid_offsets))
                


                onsets = sorted(valid_onsets)
                offsets = sorted(valid_offsets)
                
                num_intervals = len(onsets)

                # 根据min pause 调整onset和offset
                bad_onsets = []
                bad_offsets = []

                for i in range(len(onsets)-1):
                    if onsets[i+1] - offsets[i] < min_pause:
                        bad_onsets.append(onsets[i+1])
                        bad_offsets.append(offsets[i])

                onsets = [x for x in onsets if x not in bad_onsets]
                offsets = [x for x in offsets if x not in bad_offsets]

                os.makedirs(tmp_path, exist_ok=True)  # 确保临时目录存在

                audio_obj_clipped = None
                for i, (onset, offset) in enumerate(zip(onsets, offsets)):
                    if audio_obj_clipped is None:
                        audio_obj_clipped = selected_audio[onset*1000:offset*1000]
                    else:
                        audio_obj_clipped += selected_audio[onset*1000:offset*1000]
                    
                    if i != len(onsets)-1:
                        # 创建500ms空白音频
                        silence_duration = 0.2  # 500ms
                        silence_samples = int(silence_duration * selected_audio.frame_rate)
                        silence_arr = np.zeros(silence_samples, dtype=np.float32)
                        silence_audio = ReadSound(fpath=None, arr=silence_arr, duration_seconds=silence_duration, frame_rate=selected_audio.frame_rate)
                        audio_obj_clipped += silence_audio 
                # 为裁剪的音频文件生成唯一的文件名
                base_name = os.path.splitext(os.path.basename(wav_path))[0]
                clip_path = os.path.join(tmp_path, f"{base_name}_clip_{i}.wav")
                audio_obj_clipped.save(clip_path)

            
                # 转录并获取结果
                clip_result = self.model.transcribe(clip_path)
                clip_transcript = clip_result[0][0]["text_tn"]

                similarity = jellyfish.jaro_winkler_similarity(clip_transcript, standard_transcript)
                similarity = min(similarity, 0.9)
                
            else:
                similarity = 0.
                num_intervals = 0  # 可能有潜在bug，即所有num_intervals都等于0（日后待修）

            # 复制两个变量
            num_intervals_copy = num_intervals
            adjusted_params_copy = copy.deepcopy(adjusted_params)

            # 然后可以像原来一样使用复制后的变量
            return [num_intervals_copy, similarity, adjusted_params_copy]
            # result.append([num_intervals_copy, similarity, adjusted_params_copy])
        
        with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:  # 使用线程池并发处理
            result = list(tqdm(executor.map(grid_search_optimal_params, generate_param_grid(param_grid)),
                              total=len(list(generate_param_grid(param_grid))),
                              desc=f"{file_info} Locate optimal params", leave=False))
        
        max_result = max(result, key=lambda x: (x[1], x[0]))
        max_item, best_params = max_result[0], max_result[2]
        if verbose:
            print(f"[{show_elapsed_time()}] ({os.path.basename(wav_path)}) VAD best numIntervals: {max_item}, params: {best_params['onset']}")

        self.params = best_params

    def release_resources(self):
        """释放模型资源"""
        if hasattr(self, 'model') and self.model is not None:
            # 如果是直接模式，尝试释放模型
            if hasattr(self.model, 'to'):
                self.model.to('cpu')
            # 清空模型引用
            del self.model
        # 清空资源
        clear_resources()
        print("Resources released successfully")

# 导出clear_resources函数
__all__ = ['init_model', 'clear_resources']