#!/opt/homebrew/bin/python3.11
"""
Audio Taxonomy Processor (Direct AST Inference)
職責：直接調用 AST 模型進行分類，避開 HF Pipeline 的依賴問題。
"""

import numpy as np
import json
import os
import sys
import librosa
import torch

# 依賴檢查
try:
    from transformers import AutoFeatureExtractor, ASTForAudioClassification

    HAS_AST = True
except ImportError:
    print("❌ transformers not found. Run: pip install transformers")
    sys.exit(1)

# 設定
UUID = os.getenv("UUID", "384b0ff44aaaa1f1")
OUTPUT_DIR = os.getenv("MOMENTRY_OUTPUT_DIR", "./output")
AUDIO_PATH = os.path.join(OUTPUT_DIR, UUID, f"{UUID}.wav")
OUTPUT_JSON = os.path.join(OUTPUT_DIR, UUID, f"{UUID}.audio_taxonomy.json")

# 1. 標籤映射 (AudioSet -> 業務分類)
TAXONOMY_MAP = {
    "Speech": "Human/Speech",
    "Male speech, man speaking": "Human/Speech",
    "Female speech, woman speaking": "Human/Speech",
    "Conversation": "Human/Speech",
    "Laughter": "Human/Vocals",
    "Singing": "Human/Vocals",
    "Choir": "Human/Vocals",
    "Cough": "Human/Vocals",
    "Applause": "Human/Vocals",
    "Rain": "Nature/Weather",
    "Raindrop": "Nature/Weather",
    "Thunder": "Nature/Weather",
    "Wind": "Nature/Weather",
    "Ocean": "Nature/Water",
    "Stream": "Nature/Water",
    "Bird": "Nature/Flora_Fauna",
    "Dog": "Nature/Flora_Fauna",
    "Cat": "Nature/Flora_Fauna",
    "Gunshot, gunfire": "Artificial/Impact_Weapon",
    "Explosion": "Artificial/Impact_Weapon",
    "Glass shatter": "Artificial/Impact_Weapon",
    "Car": "Artificial/Transport",
    "Engine": "Artificial/Transport",
    "Siren": "Artificial/Transport",
    "Piano": "Artificial/Music",
    "Guitar": "Artificial/Music",
    "Drum": "Artificial/Music",
    "Music": "Artificial/Music",
    "Keyboard": "Artificial/Household",
    "Telephone": "Artificial/Household",
    "Door": "Artificial/Household",
}


def map_to_taxonomy(logits, model):
    """將 Logits 映射到業務分類"""
    probabilities = torch.softmax(logits, dim=-1).cpu().numpy()[0]
    # 取得 Top 5 預測
    top_indices = np.argsort(probabilities)[::-1][:5]

    events = {}
    for idx in top_indices:
        score = probabilities[idx]
        # AST 模型通常將標籤映射在 model.config.id2label
        label = model.config.id2label.get(idx, f"Class_{idx}")

        # 清洗標籤 (AST 標籤通常是 "Class X" 或實際名稱，需確認)
        # AST-finetuned-audioset 的 id2label 是 AudioSet 名稱
        mapped_cat = TAXONOMY_MAP.get(label)

        # 模糊匹配 (如果標籤不在映射表中，嘗試包含關鍵字)
        if not mapped_cat:
            lower_label = label.lower()
            if "speech" in lower_label:
                mapped_cat = "Human/Speech"
            elif "music" in lower_label:
                mapped_cat = "Artificial/Music"
            elif "gun" in lower_label or "explosion" in lower_label:
                mapped_cat = "Artificial/Impact_Weapon"
            elif "rain" in lower_label or "thunder" in lower_label:
                mapped_cat = "Nature/Weather"

        if mapped_cat and score > 0.2:
            # 只保留該類別的最高分
            if mapped_cat not in events or score > events[mapped_cat]:
                events[mapped_cat] = round(float(score), 4)
    return events


def run_audio_taxonomy(audio_path, chunk_sec=1.0, hop_sec=0.5):
    """執行分類"""
    print("🔍 Loading AST model (MIT)...")
    model_name = "MIT/ast-finetuned-audioset-10-10-0.4593"

    feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
    model = ASTForAudioClassification.from_pretrained(model_name)

    print(f"📊 Analyzing audio in {chunk_sec}s chunks (hop: {hop_sec}s)...")
    y, sr = librosa.load(audio_path, sr=16000, mono=True)
    total_dur = len(y) / sr

    results = []
    current = 0.0

    print(f"⏱️  Total duration: {total_dur:.1f}s")
    while current + chunk_sec <= total_dur:
        start_sample = int(current * sr)
        end_sample = int((current + chunk_sec) * sr)
        clip = y[start_sample:end_sample]

        # 預處理為 Tensor
        inputs = feature_extractor(clip, sampling_rate=sr, return_tensors="pt")

        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits

        taxonomy = map_to_taxonomy(logits, model)

        if taxonomy:
            results.append({"timestamp": round(current, 1), "categories": taxonomy})

        current += hop_sec
        if int(current) % 30 == 0:
            print(f"   🕒 Processed: {int(current)}s / {int(total_dur)}s", flush=True)
            # Checkpoint save (simple append/overwrite logic for safety)
            if len(results) > 0 and int(current) % 300 == 0:  # Save every 5 mins
                try:
                    temp_json = OUTPUT_JSON + ".tmp"
                    with open(temp_json, "w", encoding="utf-8") as f:
                        json.dump(
                            {"audio_taxonomy": results}, f, indent=2, ensure_ascii=False
                        )
                    # print(f"   💾 Checkpoint saved ({len(results)} events).", flush=True) # Too noisy
                except Exception:
                    pass

    return results


if __name__ == "__main__":
    if not os.path.exists(AUDIO_PATH):
        AUDIO_PATH_MP4 = os.path.join(OUTPUT_DIR, UUID, f"{UUID}.mp4")
        if not os.path.exists(AUDIO_PATH_MP4):
            AUDIO_PATH_MP4 = os.path.join(OUTPUT_DIR, UUID, f"{UUID}.mov")

        if os.path.exists(AUDIO_PATH_MP4):
            print("🎥 Extracting audio from video...")
            os.system(f"ffmpeg -y -i {AUDIO_PATH_MP4} -vn -ar 16000 -ac 1 {AUDIO_PATH}")
        else:
            print("❌ No audio/video found.")
            sys.exit(1)

    print(f"🕵️‍♂️ Starting Audio Taxonomy Classification for {UUID}...")
    events = run_audio_taxonomy(AUDIO_PATH)

    with open(OUTPUT_JSON, "w", encoding="utf-8") as f:
        json.dump({"audio_taxonomy": events}, f, indent=2, ensure_ascii=False)

    print("\n🎉 Classification Complete!")
    print(f"✅ Found {len(events)} tagged audio segments.")
    print(f"💾 Saved to {OUTPUT_JSON}")
