"""
MultiMetric Eval - 多指标翻译评测工具
"""
import os
import gc
import json
import numpy as np
import sacrebleu
import torch
from typing import Dict, List, Optional, Union
from pathlib import Path

# ==================== 配置 ====================

CACHE_PATHS = {
    "huggingface": os.path.expanduser("~/.cache/huggingface/hub"),
    "whisper": os.path.expanduser("~/.cache/whisper"),
}

# ==================== 可选依赖 ====================

try:
    import whisper
except ImportError:
    whisper = None

try:
    from bleurt import score as bleurt_score
except ImportError:
    bleurt_score = None

try:
    from comet import download_model, load_from_checkpoint
except ImportError:
    download_model = None
    load_from_checkpoint = None


# ==================== 输入加载工具 ====================

def load_hypothesis_from_file(file_path: str) -> List[str]:
    """
    从文件加载用户翻译结果
    
    支持格式:
    - .json: {"hypothesis": [...]} 或 [{"id": "x", "hypothesis": "..."}, ...]
    - .txt: 每行一句
    """
    path = Path(file_path)
    
    if not path.exists():
        raise FileNotFoundError(f"文件不存在: {file_path}")
    
    suffix = path.suffix.lower()
    
    if suffix == ".json":
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        
        if isinstance(data, dict) and "hypothesis" in data:
            return data["hypothesis"]
        
        if isinstance(data, list) and len(data) > 0:
            if isinstance(data[0], dict) and "hypothesis" in data[0]:
                return [item["hypothesis"] for item in data]
            if isinstance(data[0], str):
                return data
        
        raise ValueError("JSON 格式不正确")
    
    elif suffix == ".txt":
        with open(path, "r", encoding="utf-8") as f:
            return [line.strip() for line in f if line.strip()]
    
    else:
        raise ValueError(f"不支持的文件格式: {suffix}")


def load_audio_from_folder(folder_path: str, extensions: tuple = (".wav", ".mp3", ".flac")) -> List[str]:
    """从文件夹加载音频文件路径"""
    folder = Path(folder_path)
    
    if not folder.exists():
        raise FileNotFoundError(f"文件夹不存在: {folder_path}")
    
    audio_files = []
    for ext in extensions:
        audio_files.extend(folder.glob(f"*{ext}"))
    
    audio_files = sorted(audio_files, key=lambda x: x.stem)
    
    if not audio_files:
        raise ValueError(f"文件夹中没有音频文件: {folder_path}")
    
    return [str(f) for f in audio_files]


# ==================== 评测器 ====================

class ModelEvaluator:
    """多指标翻译评测器"""

    def __init__(
        self,
        use_comet: bool = True,
        use_bleurt: bool = False,
        use_whisper: bool = False,
        comet_model: str = "Unbabel/wmt22-comet-da",
        whisper_model: str = "medium",
        bleurt_path: Optional[str] = None,
        device: Optional[str] = None,
    ):
        # ===== 修改部分 START =====
        if device is not None:
            self.device = device
        elif torch.cuda.is_available():
            # 支持通过环境变量 CUDA_VISIBLE_DEVICES 指定 GPU
            self.device = "cuda"
        else:
            self.device = "cpu"
        
        # 显示实际使用的 GPU 信息
        if self.device.startswith("cuda"):
            if torch.cuda.is_available():
                # 如果指定了 cuda:N，提取 GPU 编号
                if ":" in self.device:
                    gpu_id = int(self.device.split(":")[1])
                    gpu_name = torch.cuda.get_device_name(gpu_id)
                    print(f"🚀 初始化评测器 (设备: {self.device} - {gpu_name})")
                else:
                    # 显示 CUDA_VISIBLE_DEVICES 环境变量（如果设置了）
                    visible = os.environ.get("CUDA_VISIBLE_DEVICES", "all")
                    gpu_name = torch.cuda.get_device_name(0)
                    print(f"🚀 初始化评测器 (设备: {self.device} - {gpu_name}, CUDA_VISIBLE_DEVICES={visible})")
            else:
                print(f"⚠️ 指定了 {self.device} 但 CUDA 不可用，回退到 CPU")
                self.device = "cpu"
                print(f"🚀 初始化评测器 (设备: cpu)")
        else:
            print(f"🚀 初始化评测器 (设备: {self.device})")
        # ===== 修改部分 END =====

        self.comet = self._load_comet(comet_model) if use_comet else None
        self.whisper_model = self._load_whisper(whisper_model) if use_whisper else None
        self.bleurt = self._load_bleurt(bleurt_path) if use_bleurt else None

        print("✅ 系统就绪！")

    # -------------------- 上下文管理器 --------------------
    
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cleanup()
        return False

    def cleanup(self):
        """释放模型显存"""
        if hasattr(self, 'comet') and self.comet is not None:
            del self.comet
            self.comet = None
        if hasattr(self, 'whisper_model') and self.whisper_model is not None:
            del self.whisper_model
            self.whisper_model = None
        if hasattr(self, 'bleurt') and self.bleurt is not None:
            del self.bleurt
            self.bleurt = None
        
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        gc.collect()
        print("🧹 已释放模型显存")

    # -------------------- 模型加载 --------------------

    def _load_comet(self, model_name: str):
        """加载 COMET 模型"""
        if not download_model:
            print("⚠️ COMET 未安装: pip install unbabel-comet")
            return None

        cache = os.path.join(CACHE_PATHS["huggingface"], f"models--{model_name.replace('/', '--')}")
        status = "[Local]" if os.path.exists(cache) else "[Online]"
        print(f"⏳ {status} 加载 COMET: {model_name}")

        model = load_from_checkpoint(download_model(model_name))
        model = model.to(self.device) if self.device == "cuda" else model
        print("✅ COMET 加载成功！")
        return model

    def _load_whisper(self, model_name: str):
        """加载 Whisper 模型"""
        if not whisper:
            print("⚠️ Whisper 未安装: pip install openai-whisper")
            return None

        cache = os.path.join(CACHE_PATHS["whisper"], f"{model_name}.pt")
        status = "[Local]" if os.path.exists(cache) else "[Online]"
        print(f"⏳ {status} 加载 Whisper: {model_name}")

        model = whisper.load_model(model_name, device=self.device)
        print("✅ Whisper 加载成功！")
        return model

    def _load_bleurt(self, path: Optional[str]):
        """加载 BLEURT 模型（强制 CPU，避免与 PyTorch 冲突）"""
        if not bleurt_score:
            print("⚠️ BLEURT 未安装: pip install bleurt")
            return None
        if not path or not os.path.exists(path):
            print(f"⚠️ BLEURT 路径无效: {path}")
            return None

        print(f"⏳ [Local] 加载 BLEURT: {path}")
        try:
            import tensorflow as tf
            tf.config.set_visible_devices([], 'GPU')
            
            scorer = bleurt_score.BleurtScorer(path)
            print("✅ BLEURT 加载成功！")
            return scorer
        except Exception as e:
            print(f"⚠️ BLEURT 加载失败: {e}")
            return None

    # -------------------- 核心内部方法 --------------------

    def transcribe(self, audio_paths: List[str]) -> List[str]:
        """语音转文字"""
        if not self.whisper_model:
            raise RuntimeError("请设置 use_whisper=True")

        print(f"🎤 ASR 转写 ({len(audio_paths)} 个文件)...")
        results = []
        for i, path in enumerate(audio_paths, 1):
            if not os.path.exists(path):
                print(f"   ⚠️ [{i}] 文件不存在: {path}")
                results.append("")
            else:
                try:
                    text = self.whisper_model.transcribe(path, fp16=(self.device == "cuda"))["text"]
                    results.append(text.strip())
                    print(f"   ✓ [{i}/{len(audio_paths)}] {os.path.basename(path)}")
                except Exception as e:
                    print(f"   ⚠️ [{i}] 转写失败: {e}")
                    results.append("")
        return results

    def _compute_metrics(
        self,
        hypothesis: List[str],
        reference: List[str],
        source: Optional[List[str]] = None,
        suffix: str = "",
    ) -> Dict[str, float]:
        """
        计算一组指标（内部方法）
        
        Args:
            hypothesis: 翻译结果
            reference: 参考译文
            source: 源文本（COMET 需要）
            suffix: 指标名后缀，如 "_ASR"
        
        Returns:
            指标字典
        """
        results = {}

        # sacreBLEU
        results[f"sacreBLEU{suffix}"] = self._safe_calc(
            lambda: sacrebleu.corpus_bleu(hypothesis, [reference]).score
        )

        # chrF++
        results[f"chrF++{suffix}"] = self._safe_calc(
            lambda: sacrebleu.corpus_chrf(hypothesis, [reference], word_order=2).score
        )

        # BLEURT
        if self.bleurt:
            results[f"BLEURT{suffix}"] = self._safe_calc(
                lambda: float(np.mean(self.bleurt.score(references=reference, candidates=hypothesis)))
            )

        # COMET
        if self.comet:
            if source:
                data = [{"src": s, "mt": h, "ref": r} for s, h, r in zip(source, hypothesis, reference)]
                gpus = 1 if self.device == "cuda" else 0
                results[f"COMET{suffix}"] = self._safe_calc(
                    lambda: self.comet.predict(data, batch_size=8, gpus=gpus).system_score
                )
            else:
                print(f"⚠️ COMET{suffix} 需要 source 参数")
                results[f"COMET{suffix}"] = -1.0

        # 四舍五入
        results = {k: round(v, 4) if v >= 0 else v for k, v in results.items()}
        return results

    # -------------------- 公开 API --------------------

    def evaluate(
        self,
        hypothesis: List[str],
        reference: List[str],
        source: Optional[List[str]] = None,
    ) -> Dict[str, float]:
        """
        纯文本评测（不涉及 ASR）
        
        Args:
            hypothesis: 用户翻译结果 (target_text)
            reference: 参考译文
            source: 源文本（COMET 需要）
        
        Returns:
            {"sacreBLEU": ..., "chrF++": ..., "COMET": ..., "BLEURT": ...}
        """
        print("📊 计算文本指标...")
        return self._compute_metrics(hypothesis, reference, source, suffix="")

    def evaluate_file(
        self,
        hypothesis_file: str,
        reference: List[str],
        source: Optional[List[str]] = None,
    ) -> Dict[str, float]:
        """从文件加载翻译结果并评测"""
        print(f"📂 加载翻译结果: {hypothesis_file}")
        hypothesis = load_hypothesis_from_file(hypothesis_file)
        print(f"   加载了 {len(hypothesis)} 条翻译")
        return self.evaluate(hypothesis, reference, source)

    def evaluate_audio_folder(
        self,
        audio_folder: str,
        reference: List[str],
        source: Optional[List[str]] = None,
    ) -> Dict[str, Union[float, List[str]]]:
        """从音频文件夹评测（纯 ASR 模式）"""
        print(f"📂 加载音频文件夹: {audio_folder}")
        audio_paths = load_audio_from_folder(audio_folder)
        print(f"   找到 {len(audio_paths)} 个音频文件")
        
        hypothesis_asr = self.transcribe(audio_paths)
        
        print("📊 计算 ASR 指标...")
        results = self._compute_metrics(hypothesis_asr, reference, source, suffix="_ASR")
        results["hypothesis_ASR"] = hypothesis_asr
        return results

    def evaluate_all(
        self,
        reference: List[str],
        source: Optional[List[str]] = None,
        target_text: Optional[Union[List[str], str]] = None,
        target_speech: Optional[str] = None,
    ) -> Dict[str, Union[float, List[str]]]:
        """
        统一评测接口（支持自定义数据）
        
        支持三种模式:
        1. 只传 target_text → 返回 sacreBLEU, chrF++, COMET, BLEURT
        2. 只传 target_speech → 返回 sacreBLEU_ASR, chrF++_ASR, COMET_ASR, BLEURT_ASR, hypothesis_ASR
        3. 同时传两者 → 返回全部指标
        
        Args:
            reference: 参考译文
            source: 源文本（COMET 需要）
            target_text: 文本翻译结果（List[str] 或文件路径）
            target_speech: 音频文件夹路径
        
        Returns:
            评测结果字典
        """
        if target_text is None and target_speech is None:
            raise ValueError("请至少提供 target_text 或 target_speech 之一")

        results = {}

        # ---- 文本评测 ----
        if target_text is not None:
            # 如果传的是文件路径，先加载
            if isinstance(target_text, str):
                print(f"📂 加载翻译文件: {target_text}")
                hyp_text = load_hypothesis_from_file(target_text)
                print(f"   加载了 {len(hyp_text)} 条翻译")
            else:
                hyp_text = target_text

            assert len(hyp_text) == len(reference), \
                f"target_text 数量 ({len(hyp_text)}) 与 reference 数量 ({len(reference)}) 不一致"

            print("📊 计算文本指标...")
            text_results = self._compute_metrics(hyp_text, reference, source, suffix="")
            results.update(text_results)
            results["hypothesis_text"] = hyp_text

        # ---- 语音评测 ----
        if target_speech is not None:
            if not self.whisper_model:
                raise RuntimeError("使用 target_speech 需要设置 use_whisper=True")

            print(f"📂 加载音频文件夹: {target_speech}")
            audio_paths = load_audio_from_folder(target_speech)
            print(f"   找到 {len(audio_paths)} 个音频文件")

            assert len(audio_paths) == len(reference), \
                f"音频文件数量 ({len(audio_paths)}) 与 reference 数量 ({len(reference)}) 不一致"

            # 如果同时有 target_text，检查数量一致
            if target_text is not None:
                hyp_text_len = len(hyp_text) if isinstance(target_text, list) else len(load_hypothesis_from_file(target_text))
                assert len(audio_paths) == hyp_text_len, \
                    f"target_text 数量 ({hyp_text_len}) 与 target_speech 数量 ({len(audio_paths)}) 不一致，同时输入时必须是同一批样本"

            hyp_asr = self.transcribe(audio_paths)

            print("📊 计算 ASR 指标...")
            asr_results = self._compute_metrics(hyp_asr, reference, source, suffix="_ASR")
            results.update(asr_results)
            results["hypothesis_ASR"] = hyp_asr

        return results

    def evaluate_dataset(
        self,
        dataset,
        target_text: Optional[Union[List[str], str]] = None,
        target_speech: Optional[str] = None,
        # 向后兼容旧参数
        hypothesis: Optional[Union[List[str], str]] = None,
        audio_folder: Optional[str] = None,
    ) -> Dict[str, Union[float, List[str]]]:
        """
        使用数据集评测（支持同时输入文本和语音）
        
        支持三种模式:
        1. 只传 target_text → 返回 sacreBLEU, chrF++, COMET, BLEURT
        2. 只传 target_speech → 返回 sacreBLEU_ASR, chrF++_ASR, COMET_ASR, BLEURT_ASR
        3. 同时传两者 → 返回全部指标
        
        Args:
            dataset: Dataset 对象
            target_text: 文本翻译结果（List[str] 或文件路径）
            target_speech: 音频文件夹路径
            hypothesis: [向后兼容] 等同于 target_text
            audio_folder: [向后兼容] 等同于 target_speech
        
        Returns:
            评测结果字典
        """
        # 向后兼容：旧参数映射到新参数
        if hypothesis is not None and target_text is None:
            target_text = hypothesis
        if audio_folder is not None and target_speech is None:
            target_speech = audio_folder

        return self.evaluate_all(
            reference=dataset.reference_texts,
            source=dataset.source_texts,
            target_text=target_text,
            target_speech=target_speech,
        )

    # -------------------- 工具方法 --------------------

    @staticmethod
    def _safe_calc(fn, default=-1.0) -> float:
        try:
            return fn()
        except:
            return default