import os
import glob
import json
import logging
import subprocess
from concurrent.futures import ProcessPoolExecutor, as_completed
import pandas as pd
from tqdm import tqdm

logging.basicConfig(level=logging.INFO)

BAK_JSONL = r"D:/dataset/audio/New folder/thaiserextract/metadata.jsonl.bak"
PARQUET_DIR = r"D:/dataset/audio/New folder/thaiser/data"
OUTPUT_AUDIO_DIR = r"D:/dataset/audio/New folder/thaiserextract/audio"
OUTPUT_JSONL = r"D:/dataset/audio/New folder/thaiserextract/metadata_clean.jsonl"

os.makedirs(OUTPUT_AUDIO_DIR, exist_ok=True)

def extract_base_id(audio_id):
    return audio_id.replace("_con_", "_").replace("_clip_", "_").replace("_middle_", "_")

def safe_json_loads(line):
    # Hack to avoid cp1252 weirdness since some lines might be broken in Windows console
    # but we are in python so utf-8 should just work if we ignore errors.
    try:
        return json.loads(line)
    except:
        return None

def load_bak_mapping():
    mapping = {}
    if os.path.exists(BAK_JSONL):
        with open(BAK_JSONL, "r", encoding="utf-8", errors="ignore") as f:
            for line in f:
                data = safe_json_loads(line)
                if data and "audio_id" in data and "text" in data:
                    base_id = extract_base_id(data["audio_id"])
                    mapping[base_id] = {
                        "text": data["text"],
                        "speaker_id": data.get("speaker_id", "Unknown")
                    }
    return mapping

def convert_bytes_to_mp3(audio_bytes, out_mp3):
    if os.path.exists(out_mp3):
        return True
    try:
        cmd = ['ffmpeg', '-y', '-f', 'flac', '-i', 'pipe:0', '-c:a', 'libmp3lame', '-q:a', '2', out_mp3]
        subprocess.run(cmd, input=audio_bytes, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        return True
    except Exception as e:
        return False

def process_parquet(parquet_file, mapping):
    results = []
    try:
        df = pd.read_parquet(parquet_file)
        # Filter matching rows
        df = df[df['assigned_emo'].notna() & (df['assigned_emo'].astype(str).str.lower() != 'none')]
        
        for _, row in df.iterrows():
            audio_id_con = row['audio_id']
            base_id = extract_base_id(audio_id_con)
            
            text = None
            speaker_id = str(row.get('actor_id', 'Unknown'))
            emo = row['assigned_emo']

            if base_id in mapping:
                text = mapping[base_id]["text"]
                speaker_id = mapping[base_id]["speaker_id"]
            elif pd.notna(row.get('script_sent')):
                text = str(row['script_sent'])
            
            if not text or len(text.strip()) == 0:
                continue

            mics = [("mic_con", "_con_"), ("mic_clip", "_clip_"), ("mic_middle", "_middle_")]
            # Construct the base name parts. Usually "s001_actor001_..."
            # We want to insert the mic infix properly.
            # Example base_id: s001_actor001_impro1_1.flac
            # So replace the first '_' with '_con_' -> s001_con_actor001_impro1_1.flac
            for col, infix in mics:
                if col in row and pd.notna(row[col]) and isinstance(row[col], dict) and 'bytes' in row[col]:
                    first_underscore = base_id.find('_')
                    if first_underscore != -1:
                        new_audio_id = base_id[:first_underscore] + infix + base_id[first_underscore+1:]
                    else:
                        new_audio_id = infix.strip('_') + "_" + base_id
                        
                    out_mp3_name = new_audio_id + ".mp3"
                    out_mp3_path = os.path.join(OUTPUT_AUDIO_DIR, out_mp3_name)
                    
                    success = convert_bytes_to_mp3(row[col]['bytes'], out_mp3_path)
                    if success:
                        results.append({
                            "audio_id": new_audio_id,
                            "file_path": out_mp3_path.replace("\\", "/"),
                            "text": text,
                            "speaker_id": speaker_id,
                            "emotion": emo,
                            "path": out_mp3_path.replace("\\", "/")
                        })
    except Exception as e:
        print(f"Error processing {parquet_file}: {e}")
    return results

def main():
    print("Loading mapping from metadata.jsonl.bak...")
    mapping = load_bak_mapping()
    print(f"Loaded {len(mapping)} verified texts.")
    
    parquet_files = glob.glob(os.path.join(PARQUET_DIR, "*.parquet"))
    print(f"Found {len(parquet_files)} parquet files to process.")
    
    all_results = []
    
    # Process sequentially for safety (ffmpeg multi processing can hit cpu limits if not careful, 
    # but let's use a small process pool 4 workers).
    with ProcessPoolExecutor(max_workers=4) as executor:
        futures = {executor.submit(process_parquet, pf, mapping): pf for pf in parquet_files}
        
        for future in tqdm(as_completed(futures), total=len(futures), desc="Extracting Audio"):
            res = future.result()
            all_results.extend(res)
            
    print(f"Extraction complete! Total multi-mic segments: {len(all_results)}")
    
    print("Writing metadata_clean.jsonl...")
    with open(OUTPUT_JSONL, "w", encoding="utf-8") as f:
        for r in all_results:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print("Done!")

if __name__ == "__main__":
    main()