"""
MUSIC DOSSIER AI — CLAP + Essentia Mood/Key API Server
Hugging Face Spaces (Docker) 用
v9: MTG-Jamendo 56クラスムード + 3アルゴリズムキー多数決
"""

import os
import base64
import tempfile
import subprocess
import numpy as np
import librosa
import torch
import httpx
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from typing import List, Optional
from transformers import ClapModel, ClapProcessor

# ============================================================
# モデルディレクトリ
# ============================================================
MODEL_DIR = '/app/essentia_models'
os.makedirs(MODEL_DIR, exist_ok=True)

MOOD_MODEL_URLS = {
    'mood_happy':      'https://essentia.upf.edu/models/classification-heads/mood_happy/mood_happy-msd-musicnn-1.pb',
    'mood_sad':        'https://essentia.upf.edu/models/classification-heads/mood_sad/mood_sad-msd-musicnn-1.pb',
    'mood_aggressive': 'https://essentia.upf.edu/models/classification-heads/mood_aggressive/mood_aggressive-msd-musicnn-1.pb',
    'mood_relaxed':    'https://essentia.upf.edu/models/classification-heads/mood_relaxed/mood_relaxed-msd-musicnn-1.pb',
    'mood_party':      'https://essentia.upf.edu/models/classification-heads/mood_party/mood_party-msd-musicnn-1.pb',
    'msd_musicnn':     'https://essentia.upf.edu/models/feature-extractors/musicnn/msd-musicnn-1.pb',
    # v9追加: MTG-Jamendo 56クラスムード
    'discogs_effnet':  'https://essentia.upf.edu/models/feature-extractors/discogs-effnet/discogs-effnet-bs64-1.pb',
    'mtg_jamendo':     'https://essentia.upf.edu/models/classification-heads/mtg_jamendo_moodtheme/mtg_jamendo_moodtheme-discogs-effnet-1.pb',
}

# ============================================================
# 起動時にモデルをダウンロード（安全版）
# ============================================================
def download_model(name, url, path):
    """モデルをダウンロードし、成否とファイルサイズを検証する"""
    result = subprocess.run(
        ['wget', '-q', '--tries=3', '--timeout=60', '-O', path, url]
    )
    if result.returncode != 0:
        print(f'  ⚠️ {name} ダウンロード失敗 (exit={result.returncode})')
        if os.path.exists(path):
            os.remove(path)
        return False
    size = os.path.getsize(path) if os.path.exists(path) else 0
    if size < 10000:
        print(f'  ⚠️ {name} ファイルが空または破損 ({size}bytes) → 削除')
        os.remove(path)
        return False
    print(f'  ✅ {name} ダウンロード完了 ({size//1024}KB)')
    return True

import time

print('🔄 Essentiaムードモデルをダウンロード中...')
for name, url in MOOD_MODEL_URLS.items():
    path = f'{MODEL_DIR}/{name}.pb'
    if os.path.exists(path):
        size = os.path.getsize(path)
        if size > 10000:
            print(f'  キャッシュ済み: {name} ({size//1024}KB)')
        else:
            print(f'  ⚠️ {name} キャッシュが破損 → 再ダウンロード')
            os.remove(path)
            download_model(name, url, path)
            time.sleep(3)  # レート制限対策
    else:
        print(f'  ダウンロード: {name}...')
        success = download_model(name, url, path)
        if success:
            time.sleep(3)  # 連続ダウンロードによるレート制限を回避
        else:
            time.sleep(10)  # 失敗時は長めに待つ

# ============================================================
# モデルのロード
# ============================================================
print('🔄 Essentiaムードモデルをロード中...')
import essentia.standard as es

# ── 基本ムードモデル（必須） ──
msd_path = f'{MODEL_DIR}/msd_musicnn.pb'
if not os.path.exists(msd_path) or os.path.getsize(msd_path) < 10000:
    raise RuntimeError(f'msd_musicnn.pb が存在しないか破損しています。URLを確認してください。')

embedding_model = es.TensorflowPredictMusiCNN(
    graphFilename=msd_path,
    output='model/dense/BiasAdd'
)

mood_models = {}
for name in ['mood_happy', 'mood_sad', 'mood_aggressive', 'mood_relaxed', 'mood_party']:
    path = f'{MODEL_DIR}/{name}.pb'
    if os.path.exists(path) and os.path.getsize(path) > 10000:
        try:
            mood_models[name] = es.TensorflowPredict2D(
                graphFilename=path,
                output='model/Softmax'
            )
            print(f'  ✅ {name} ロード完了')
        except Exception as e:
            print(f'  ⚠️ {name} ロード失敗（スキップ）: {e}')
    else:
        print(f'  ⚠️ {name} ファイルなし・スキップ')

# ── MTG-Jamendo（v9追加・オプション） ──
effnet_embedding_model = None
mtg_jamendo_model = None

effnet_path = f'{MODEL_DIR}/discogs_effnet.pb'
mtg_path = f'{MODEL_DIR}/mtg_jamendo.pb'

if (os.path.exists(effnet_path) and os.path.getsize(effnet_path) > 10000 and
    os.path.exists(mtg_path) and os.path.getsize(mtg_path) > 10000):
    try:
        effnet_embedding_model = es.TensorflowPredictEffnetDiscogs(
            graphFilename=effnet_path,
            output='PartitionedCall:1'
        )
        mtg_jamendo_model = es.TensorflowPredict2D(
            graphFilename=mtg_path,
            output='model/Sigmoid',
            input='model/Placeholder'
        )
        print('  ✅ MTG-Jamendo 56クラスモデル ロード完了')
    except Exception as e:
        print(f'  ⚠️ MTG-Jamendoモデルのロード失敗（スキップ）: {e}')
        effnet_embedding_model = None
        mtg_jamendo_model = None
else:
    print('  ⚠️ MTG-Jamendo / discogs_effnet ファイルなし・スキップ')

MTG_JAMENDO_LABELS = [
    'action', 'adventure', 'advertising', 'background', 'ballad',
    'calm', 'children', 'christmas', 'commercial', 'cool',
    'corporate', 'dark', 'deep', 'documentary', 'drama',
    'dramatic', 'dream', 'emotional', 'energetic', 'epic',
    'fast', 'film', 'fun', 'funny', 'game',
    'groovy', 'happy', 'heavy', 'holiday', 'hopeful',
    'inspiring', 'love', 'meditative', 'melancholic', 'melodic',
    'motivational', 'movie', 'nature', 'party', 'positive',
    'powerful', 'relaxing', 'retro', 'romantic', 'sad',
    'sexy', 'slow', 'soft', 'soundscape', 'space',
    'sport', 'summer', 'trailer', 'travel', 'upbeat',
    'uplifting'
]

KEY_PROFILES = ['temperley', 'krumhansl', 'edma']

print('✅ Essentiaムードモデルのロード完了')

print('🔄 CLAPモデルをロード中...')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'使用デバイス: {device}')
clap_processor = ClapProcessor.from_pretrained('laion/clap-htsat-unfused')
clap_model = ClapModel.from_pretrained('laion/clap-htsat-unfused').to(device)
print('✅ CLAPモデルのロード完了')

# ============================================================
# FastAPI
# ============================================================
app = FastAPI(title='MUSIC DOSSIER AI API')

app.add_middleware(
    CORSMiddleware,
    allow_origins=['*'],
    allow_methods=['*'],
    allow_headers=['*'],
)

class AnalyzeRequest(BaseModel):
    audio_base64: str
    audio_format: str
    tags: List[str]
    top_k: int = 30

@app.get('/')
def health_check():
    return {
        'status': 'ok',
        'models': 'CLAP + Essentia Mood',
        'device': device,
        'mtg_jamendo_available': mtg_jamendo_model is not None,
        'mood_models_loaded': list(mood_models.keys()),
    }

@app.post('/analyze')
async def analyze_audio(req: AnalyzeRequest):
    audio_bytes = base64.b64decode(req.audio_base64)
    suffix = f'.{req.audio_format}'

    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(audio_bytes)
        tmp_path = tmp.name

    try:
        # ── CLAP解析 ──
        audio_array, sr = librosa.load(tmp_path, sr=48000, mono=True)
        inputs = clap_processor(
            audios=audio_array, return_tensors='pt', sampling_rate=48000
        ).to(device)
        text_inputs = clap_processor(
            text=req.tags, return_tensors='pt', padding=True
        ).to(device)

        with torch.no_grad():
            audio_emb = clap_model.get_audio_features(**inputs)
            text_emb  = clap_model.get_text_features(**text_inputs)
            audio_emb = audio_emb / audio_emb.norm(dim=-1, keepdim=True)
            text_emb  = text_emb  / text_emb.norm(dim=-1, keepdim=True)
            sims = (audio_emb @ text_emb.T)[0].cpu().numpy()

        clap_scores = {tag: float(s) for tag, s in zip(req.tags, sims)}
        top_tags = sorted(clap_scores, key=clap_scores.get, reverse=True)[:req.top_k]

        # ── Essentia 基本ムード解析 ──
        audio_16k, _ = librosa.load(tmp_path, sr=16000, mono=True)
        audio_es = audio_16k.astype(np.float32)
        embeddings = embedding_model(audio_es)

        mood_label_map = {
            'mood_happy': 'happy', 'mood_sad': 'sad',
            'mood_aggressive': 'aggressive', 'mood_relaxed': 'relaxed',
            'mood_party': 'party',
        }
        mood_scores = {}
        for model_name, label in mood_label_map.items():
            if model_name in mood_models:
                preds = mood_models[model_name](embeddings)
                mood_scores[label] = round(float(np.mean(preds[:, 1])), 4)

        # ── MTG-Jamendo 56クラス解析（オプション）──
        mood_scores_advanced = {}
        if effnet_embedding_model is not None and mtg_jamendo_model is not None:
            try:
                audio_44k, _ = librosa.load(tmp_path, sr=44100, mono=True)
                effnet_embeddings = effnet_embedding_model(audio_44k.astype(np.float32))
                mtg_preds = mtg_jamendo_model(effnet_embeddings)
                mtg_avg = np.mean(mtg_preds, axis=0)
                mood_scores_advanced = {
                    label: round(float(score), 4)
                    for label, score in zip(MTG_JAMENDO_LABELS, mtg_avg)
                }
            except Exception as e:
                print(f'[MTG-Jamendo] 解析エラー（スキップ）: {e}')

        return {
            'clap_scores': clap_scores,
            'top_tags': top_tags,
            'mood_scores': mood_scores,
            'mood_scores_advanced': mood_scores_advanced,
            'device': device
        }

    finally:
        os.unlink(tmp_path)

# ============================================================
# v9: キー判定エンドポイント（3アルゴリズム多数決）
# ============================================================
class KeyAnalyzeRequest(BaseModel):
    audio_base64: str
    audio_format: str

@app.post('/analyze/key')
async def analyze_key(req: KeyAnalyzeRequest):
    audio_bytes = base64.b64decode(req.audio_base64)
    suffix = f'.{req.audio_format}'

    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(audio_bytes)
        tmp_path = tmp.name

    try:
        audio_vec = es.MonoLoader(filename=tmp_path, sampleRate=44100)()
        results = []
        for profile in KEY_PROFILES:
            try:
                key, scale, strength = es.KeyExtractor(profileType=profile)(audio_vec)
                results.append({
                    'profile': profile, 'key': key, 'scale': scale,
                    'strength': round(float(strength), 4),
                    'result': f'{key} {scale}'
                })
            except Exception as e:
                print(f'[KeyExtractor:{profile}] エラー: {e}')

        from collections import Counter
        votes = Counter(r['result'] for r in results)
        winner = votes.most_common(1)[0][0]
        method = 'majority_vote'
        if votes.most_common(1)[0][1] == 1 and len(results) > 1:
            winner = max(results, key=lambda r: r['strength'])['result']
            method = 'strength_max'

        return {'key': winner, 'details': results, 'method': method, 'vote_counts': dict(votes)}

    finally:
        os.unlink(tmp_path)

# ============================================================
# ProducerAI (TTAPI) プロキシ
# ============================================================
TTAPI_BASE = 'https://api.ttapi.io'

class ProducerMusicRequest(BaseModel):
    sound_prompt: str
    lyrics_text: Optional[str] = None
    instrumental: Optional[bool] = False
    mv: Optional[str] = 'FUZZ-2.0'
    title: Optional[str] = None
    lyrics_strength: Optional[float] = 0.5
    weirdness: Optional[float] = 0.5

class ProducerFetchRequest(BaseModel):
    jobId: str

@app.post('/proxy/producer/music')
async def proxy_producer_music(req: ProducerMusicRequest, request: Request):
    tt_api_key = request.headers.get('TT-API-KEY', '')
    if not tt_api_key:
        return JSONResponse({'status': 'ERROR', 'message': 'TT-API-KEY header missing'}, status_code=400)
    async with httpx.AsyncClient(timeout=30) as client:
        res = await client.post(
            f'{TTAPI_BASE}/producer/v1/music',
            headers={'Content-Type': 'application/json', 'TT-API-KEY': tt_api_key},
            json=req.dict(exclude_none=True)
        )
    return JSONResponse(res.json(), status_code=res.status_code)

@app.post('/proxy/producer/fetch')
async def proxy_producer_fetch(req: ProducerFetchRequest, request: Request):
    tt_api_key = request.headers.get('TT-API-KEY', '')
    if not tt_api_key:
        return JSONResponse({'status': 'ERROR', 'message': 'TT-API-KEY header missing'}, status_code=400)
    async with httpx.AsyncClient(timeout=30) as client:
        res = await client.get(
            f'{TTAPI_BASE}/producer/v1/fetch',
            headers={'TT-API-KEY': tt_api_key},
            params={'jobId': req.jobId}
        )
        if res.status_code == 405:
            res = await client.post(
                f'{TTAPI_BASE}/producer/v1/fetch',
                headers={'Content-Type': 'application/json', 'TT-API-KEY': tt_api_key},
                json=req.dict()
            )
    return JSONResponse(res.json(), status_code=res.status_code)

# ============================================================
# 起動
# ============================================================
if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host='0.0.0.0', port=7860)