import os
import sys
import pathlib
import numpy as np
import torch
import torchaudio
import torchaudio.functional as F
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from tempfile import NamedTemporaryFile
from pathlib import Path

# 尝试导入 FBank
try:
    from speakerlab.process.processor import FBank
except ImportError:
    sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
    from speakerlab.process.processor import FBank

from speakerlab.utils.builder import dynamic_import

# ========== 固定使用你的模型 ==========
YOUR_MODEL_ID = "iic/speech_campplus_sv_zh-cn_16k-common"
MODEL_DIR_NAME = "iic/speech_campplus_sv_zh-cn_16k-common"  # 本地文件夹名
WEIGHT_FILENAME = "campplus_cn_common.bin"

# 模型结构配置
CAMPPLUS_COMMON = {
    'obj': 'speakerlab.models.campplus.DTDNN.CAMPPlus',
    'args': {
        'feat_dim': 80,
        'embedding_size': 192,
    },
}

# 默认本地模型根目录（可修改）
DEFAULT_MODEL_ROOT = "./models"

# ========== 全局缓存 ==========
_embedding_model = None
_feature_extractor = None
_device = None

def load_model_once(model_root: str = DEFAULT_MODEL_ROOT):
    global _embedding_model, _feature_extractor, _device
    if _embedding_model is not None:
        return _embedding_model, _feature_extractor, _device

    model_dir = Path(model_root) / MODEL_DIR_NAME
    weight_path = model_dir / WEIGHT_FILENAME

    if not weight_path.exists():
        raise RuntimeError(
            f"Model weight file not found: {weight_path}\n"
            f"Please download the model and place it under `{model_dir}`"
        )

    print(f"[INFO] Loading model from {weight_path}")
    pretrained_state = torch.load(weight_path, map_location='cpu')

    _device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model_class = dynamic_import(CAMPPLUS_COMMON['obj'])
    _embedding_model = model_class(**CAMPPLUS_COMMON['args'])
    _embedding_model.load_state_dict(pretrained_state)
    _embedding_model.to(_device)
    _embedding_model.eval()

    _feature_extractor = FBank(80, sample_rate=16000, mean_nor=True)

    print(f"[INFO] Model loaded on {_device}")
    return _embedding_model, _feature_extractor, _device

def load_wav(wav_file, obj_fs=16000):
    wav, fs = torchaudio.load(wav_file)
    if fs != obj_fs:
        # 使用 torchaudio 内置重采样（无需 sox）
        wav = F.resample(wav, orig_freq=fs, new_freq=obj_fs)
        fs = obj_fs
    if wav.shape[0] > 1:
        wav = wav[0, :].unsqueeze(0)  # 取单通道
    return wav
# def load_wav(wav_file, obj_fs=16000):
#     wav, fs = torchaudio.load(wav_file)
#     if fs != obj_fs:
#         wav, fs = torchaudio.sox_effects.apply_effects_tensor(
#             wav, fs, effects=[['rate', str(obj_fs)]]
#         )
#     if wav.shape[0] > 1:
#         wav = wav[0, :].unsqueeze(0)
#     return wav

def compute_embedding(wav_path: str):
    global _embedding_model, _feature_extractor, _device
    if _embedding_model is None:
        load_model_once()

    wav = load_wav(wav_path)
    feat = _feature_extractor(wav).unsqueeze(0).to(_device)
    with torch.no_grad():
        embedding = _embedding_model(feat).detach().squeeze(0).cpu().numpy()
    return embedding.tolist()

# ========== FastAPI App ==========
app = FastAPI(title="Speaker Embedding Extractor (Your Model)", version="1.2")

@app.post("/extract_embedding")
async def extract_embedding(wav_file: UploadFile = File(...)):
    if not wav_file.filename.endswith(".wav"):
        raise HTTPException(status_code=400, detail="Only .wav files are supported.")

    with NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
        tmp.write(await wav_file.read())
        tmp_path = tmp.name

    try:
        embedding = compute_embedding(tmp_path)
        return JSONResponse({"embedding": embedding})
    finally:
        os.unlink(tmp_path)

@app.post("/compare_similarity")
async def compare_similarity(
    wav_file1: UploadFile = File(...),
    wav_file2: UploadFile = File(...)
):
    for f in [wav_file1, wav_file2]:
        if not f.filename.endswith(".wav"):
            raise HTTPException(status_code=400, detail="Both files must be .wav")

    tmp_paths = []
    try:
        for wav_file in [wav_file1, wav_file2]:
            with NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
                tmp.write(await wav_file.read())
                tmp_paths.append(tmp.name)

        emb1 = torch.tensor(compute_embedding(tmp_paths[0]))
        emb2 = torch.tensor(compute_embedding(tmp_paths[1]))

        similarity = torch.nn.CosineSimilarity(dim=-1, eps=1e-6)
        score = similarity(emb1.unsqueeze(0), emb2.unsqueeze(0)).item()

        return JSONResponse({
            "similarity_score": round(score, 4),
            "embedding1": emb1.tolist(),
            "embedding2": emb2.tolist()
        })
    finally:
        for p in tmp_paths:
            if os.path.exists(p):
                os.unlink(p)

# 可选：启动时预加载模型
@app.on_event("startup")
async def startup_event():
    try:
        load_model_once()
    except Exception as e:
        print(f"[ERROR] Failed to load model on startup: {e}")