"""
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"),
}

for var in ["HF_DATASETS_OFFLINE", "TRANSFORMERS_OFFLINE"]:
    os.environ.pop(var, None)

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

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,
    ):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        print(f"🚀 初始化评测器 (设备: {self.device})")

        self.comet = self._load_comet(comet_model) if use_comet else None
        self.whisper = 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') and self.whisper is not None:
            del self.whisper
            self.whisper = 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 模型"""
        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:
            # 强制 BLEURT 使用 CPU（避免与 PyTorch 抢显存）
            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:
            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}] 文件不存在")
                results.append("")
            else:
                try:
                    text = self.whisper.transcribe(path, fp16=(self.device == "cuda"))["text"]
                    results.append(text.strip())
                    print(f"   ✓ [{i}/{len(audio_paths)}] {os.path.basename(path)}")
                except:
                    results.append("")
        return results

    def evaluate(
        self,
        hypothesis: List[str],
        reference: List[str],
        source: Optional[List[str]] = None,
    ) -> Dict[str, float]:
        """计算评测指标"""
        print("📊 计算指标...")

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

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

        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["COMET"] = self._safe_calc(lambda: self.comet.predict(data, batch_size=8, gpus=gpus).system_score)
            else:
                print("⚠️ COMET 需要 source 参数")
                results["COMET"] = -1.0

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

    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]]]:
        """从文件夹加载音频并评测"""
        print(f"📂 加载音频文件夹: {audio_folder}")
        audio_paths = load_audio_from_folder(audio_folder)
        print(f"   找到 {len(audio_paths)} 个音频文件")
        
        hypothesis = self.transcribe(audio_paths)
        results = self.evaluate(hypothesis, reference, source)
        results["hypothesis"] = hypothesis
        return results

    def evaluate_dataset(
        self,
        dataset,
        hypothesis: Optional[Union[List[str], str]] = None,
        audio_folder: Optional[str] = None,
    ) -> Dict[str, Union[float, List[str]]]:
        """使用数据集评测"""
        if hypothesis:
            if isinstance(hypothesis, str):
                hyp_list = load_hypothesis_from_file(hypothesis)
            else:
                hyp_list = hypothesis
            
            results = self.evaluate(hyp_list, dataset.reference_texts, dataset.source_texts)
            results["hypothesis"] = hyp_list
        
        elif audio_folder:
            results = self.evaluate_audio_folder(
                audio_folder, 
                dataset.reference_texts, 
                dataset.source_texts
            )
        
        else:
            raise ValueError("请提供 hypothesis（列表或文件路径）或 audio_folder")
        
        return results

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

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