# -*- coding: utf-8 -*-
"""VoxCeleb_Extraction_Colab.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/xxxx

# 🚀 VoxCeleb Audio Downloader (Robust Version)

This script:
1. Downloads the full VoxCeleb1 Dev dataset (4 parts) manually (bypassing library issues).
2. Concatenates them into a valid ZIP.
3. Downloads metadata to find Indian celebrities.
4. Extracts ONLY the Indian celebrity audio.
5. Zips the result for you to download.
"""

# 1. Install Dependencies
!pip install pandas

import os
import shutil
import zipfile
import json
import pandas as pd
from pathlib import Path

# Config
BASE_URL = "https://huggingface.co/datasets/ProgramComputer/voxceleb/resolve/main/vox1"
FILES = [
    "vox1_dev_wav_partaa",
    "vox1_dev_wav_partab",
    "vox1_dev_wav_partac",
    "vox1_dev_wav_partad"
]
META_URL = "https://mm.kaist.ac.kr/datasets/voxceleb/meta/vox1_meta.csv" # Official meta
# Alternative meta: https://huggingface.co/datasets/ProgramComputer/voxceleb/raw/main/vox1/vox1_meta.csv

WORK_DIR = Path("/content/vox_work")
OUTPUT_DIR = Path("/content/voxceleb_indian_audio")
WORK_DIR.mkdir(parents=True, exist_ok=True)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# 2. Download Metadata
print("⬇️ Downloading Metadata...")
!wget -q --no-check-certificate {META_URL} -O {WORK_DIR}/vox1_meta.csv

# 3. Filter for Indian Celebrities
print("🔍 Filtering Metadata...")
try:
    df = pd.read_csv(f"{WORK_DIR}/vox1_meta.csv", sep="\t")
    # Columns: "VoxCeleb1 ID", "VGGFace1 ID", "Gender", "Nationality", "Set"
    
    # Normalize nationality
    indian_celebs = df[df['Nationality'].str.lower().isin(['india', 'indian'])]
    
    # FIX: Explicitly get the ID column
    # The column name is usually "VoxCeleb1 ID"
    if "VoxCeleb1 ID" in df.columns:
        target_ids = set(indian_celebs["VoxCeleb1 ID"].astype(str).tolist())
    else:
        # Fallback: assume first column
        target_ids = set(indian_celebs.iloc[:, 0].astype(str).tolist())
    
    # LIMIT to 25
    target_ids = set(list(target_ids)[:25])

    print(f"🇮🇳 Found {len(target_ids)} Indian celebrities in metadata (Limited to 25).")
    print(f"   Sample IDs: {list(target_ids)[:5]}")
    
    # Create mapping ID -> Name
    if "VGGFace1 ID" in df.columns:
        id_to_name = dict(zip(indian_celebs["VoxCeleb1 ID"], indian_celebs["VGGFace1 ID"]))
    else:
        id_to_name = {}

except Exception as e:
    print(f"❌ Metadata Error: {e}")
    target_ids = set()

# 4. Download and Concatenate Archives
print("⬇️ Downloading Audio Archives (4 Parts) - This uses blazing fast Google Cloud...")
for f in FILES:
    url = f"{BASE_URL}/{f}"
    print(f"   Getting {f}...")
    !wget -q {url} -O {WORK_DIR}/{f}

print("🔗 Concatenating parts...")
!cat {WORK_DIR}/vox1_dev_wav_part* > {WORK_DIR}/full_vox1.zip

# 5. Extract ONLY Target Files
print(f"📦 Unzipping and extracting matching audio...")

try:
    with zipfile.ZipFile(f"{WORK_DIR}/full_vox1.zip", 'r') as z:
        file_list = z.namelist()
        extracted_count = 0
        
        for f in file_list:
            # Structure: wav/id10001/1zn.../00001.wav
            parts = f.split('/')
            if len(parts) > 2:
                s_id = parts[1] # id10001
                
                if s_id in target_ids:
                    # It's an Indian celeb!
                    name = id_to_name.get(s_id, s_id)
                    
                    # Create person dir: Name_ID
                    p_dir = OUTPUT_DIR / f"{name}_{s_id}"
                    p_dir.mkdir(exist_ok=True)
                    
                    # Target file path
                    # Limit to 10 clips per person
                    current_clips = list(p_dir.glob("*.wav"))
                    if len(current_clips) < 10:
                        if f.endswith('.wav'):
                            # Extract
                            target_path = p_dir / Path(f).name
                            with z.open(f) as source, open(target_path, "wb") as dest:
                                shutil.copyfileobj(source, dest)
                            extracted_count += 1
                    
                    # Create Metadata
                    meta = {
                        "name": name.replace("_", " "),
                        "relation": "Celebrity",
                        "nationality": "India",
                        "id": s_id
                    }
                    with open(p_dir / "metadata.json", "w") as f:
                        json.dump(meta, f)
                            
        print(f"✅ Extracted {extracted_count} total clips.")

except zipfile.BadZipFile:
    print("❌ Critical: The constructed zip file is invalid. Download might have failed.")
except Exception as e:
    print(f"❌ Error during extraction: {e}")

# 6. Zip Result
print("🗜️ Zipping final output...")
shutil.make_archive("/content/indian_celebs_audio", 'zip', OUTPUT_DIR)
print("✨ Done! Download 'indian_celebs_audio.zip' (approx ~50-100MB) from the sidebar.")