import asyncio
import aiohttp
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
import soundfile as sf
import requests
import os
from tqdm import tqdm
import pandas as pd
import json
import pyarrow as pa
import pyarrow.parquet as pq
import numpy as np
from huggingface_hub import HfApi, create_repo
from datasets import load_dataset, Audio, Dataset
import time
from functools import partial
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Configuration for 192 cores
NUM_CORES = 192
BATCH_SIZE = 32  # Increased batch size for better throughput
MAX_CONCURRENT_REQUESTS = 48  # 192/4 for optimal concurrency
CHUNK_SIZE = 1000  # Process data in chunks to manage memory

# Load the dataset without audio decoding first
print("Loading dataset...")
ds = load_dataset(
    "Ashegh-Sad-Warrior/Persian_Common_Voice_17_0",
    split="validated",
    streaming=False
)

# Now cast to audio after loading
print("Casting to audio...")
ds = ds.cast_column("audio", Audio(sampling_rate=16000))

output_dir = "confirmed_dataset"
os.makedirs(output_dir, exist_ok=True)

API_URL = "http://localhost:5000/batch_confirm"

# Hugging Face configuration
HF_DATASET_NAME = "dpr2000/persian-cv17-confirmed"
HF_PRIVATE = True

def save_flac(audio_array, path):
    """Save audio array as FLAC file"""
    sf.write(path, audio_array, 16000, format="FLAC")

def process_audio_chunk(audio_data):
    """Process a single audio item - designed for multiprocessing"""
    audio, sentence = audio_data
    flac_path = f"temp_{hash(audio.tobytes())}.flac"
    save_flac(audio["array"], flac_path)
    return {
        'flac_path': flac_path,
        'sentence': sentence,
        'audio_array': audio["array"]
    }

async def send_batch_request(session, batch_data, batch_id):
    """Send a single batch request asynchronously"""
    files = {}
    references = []
    temp_flacs = []
    audio_arrays = []
    
    for j, item in enumerate(batch_data):
        files[f"audio{j}"] = open(item['flac_path'], "rb")
        references.append(item['sentence'])
        temp_flacs.append(item['flac_path'])
        audio_arrays.append(item['audio_array'])
    
    data = {"references": json.dumps(references)}
    
    try:
        async with session.post(API_URL, data=data, files=files, timeout=aiohttp.ClientTimeout(total=120)) as response:
            if response.status == 200:
                resp_json = await response.json()
                if "results" in resp_json:
                    results = resp_json["results"]
                else:
                    logger.warning(f"Batch {batch_id} failed: 'results' key missing")
                    results = [None] * len(references)
            else:
                logger.error(f"Batch {batch_id} failed: HTTP {response.status}")
                results = [None] * len(references)
    except Exception as e:
        logger.error(f"Batch {batch_id} failed: {e}")
        results = [None] * len(references)
    finally:
        # Clean up files
        for f in files.values():
            f.close()
        for flac_path in temp_flacs:
            try:
                os.remove(flac_path)
            except:
                pass
    
    # Process results
    confirmed_items = []
    for j, result in enumerate(results):
        if result and result.get("confirmed"):
            confirmed_items.append({
                "audio": audio_arrays[j], 
                "transcription": references[j]
            })
    
    return confirmed_items

async def process_dataset_async():
    """Main async processing function"""
    confirmed = []
    
    # Prepare all audio data first using multiprocessing
    print("Preparing audio data with multiprocessing...")
    audio_data = [(ds[i]["audio"], ds[i]["sentence"]) for i in range(len(ds))]
    
    # Use ProcessPoolExecutor for CPU-intensive audio processing
    with ProcessPoolExecutor(max_workers=NUM_CORES) as executor:
        processed_audio = list(tqdm(
            executor.map(process_audio_chunk, audio_data),
            total=len(audio_data),
            desc="Processing audio files"
        ))
    
    # Create batches
    batches = []
    for i in range(0, len(processed_audio), BATCH_SIZE):
        batch = processed_audio[i:i+BATCH_SIZE]
        batches.append((batch, i // BATCH_SIZE))
    
    print(f"Processing {len(batches)} batches with {MAX_CONCURRENT_REQUESTS} concurrent requests...")
    
    # Process batches asynchronously
    async with aiohttp.ClientSession(
        connector=aiohttp.TCPConnector(limit=MAX_CONCURRENT_REQUESTS),
        timeout=aiohttp.ClientTimeout(total=300)
    ) as session:
        tasks = []
        for batch_data, batch_id in batches:
            task = send_batch_request(session, batch_data, batch_id)
            tasks.append(task)
        
        # Process in chunks to avoid overwhelming the system
        chunk_size = MAX_CONCURRENT_REQUESTS
        for i in range(0, len(tasks), chunk_size):
            chunk_tasks = tasks[i:i+chunk_size]
            results = await asyncio.gather(*chunk_tasks, return_exceptions=True)
            
            for result in results:
                if isinstance(result, Exception):
                    logger.error(f"Task failed: {result}")
                else:
                    confirmed.extend(result)
            
            print(f"Processed {min(i+chunk_size, len(tasks))}/{len(tasks)} batches, confirmed: {len(confirmed)}")
    
    return confirmed

def save_confirmed_data_parallel(confirmed):
    """Save confirmed data using parallel processing"""
    if not confirmed:
        print("❌ No confirmed samples to save")
        return
    
    print(f"\n🔄 Saving {len(confirmed)} confirmed samples...")
    
    def extract_minimal(example):
        """Convert audio to int16 format"""
        audio_float32 = np.array(example["audio"], dtype=np.float32)
        audio_float32 = np.clip(audio_float32, -1.0, 1.0)
        audio_int16 = (audio_float32 * 32767).astype(np.int16)
        return {
            "audio": audio_int16.tobytes(),
            "text": example["transcription"]
        }
    
    # Create dataset from confirmed samples
    confirmed_dataset = Dataset.from_list(confirmed)
    confirmed_dataset = confirmed_dataset.map(
        extract_minimal, 
        remove_columns=confirmed_dataset.column_names,
        num_proc=NUM_CORES  # Use all cores for dataset processing
    )
    
    # Optimize sharding for parallel writing
    num_shards = min(50, len(confirmed))  # More shards for better parallelization
    shard_size = len(confirmed_dataset) // num_shards + 1
    
    def write_shard(shard_info):
        """Write a single shard - designed for multiprocessing"""
        i, start, end = shard_info
        if start >= len(confirmed_dataset):
            return None
        
        shard = confirmed_dataset.select(range(start, end))
        table = pa.Table.from_pandas(shard.to_pandas())
        
        shard_path = os.path.join(output_dir, f"confirmed_shard_{i:03}.parquet")
        
        pq.write_table(
            table,
            shard_path,
            compression="zstd",
            compression_level=22,
            use_dictionary=True,
            version="2.6"
        )
        
        return f"Shard {i+1}: {len(shard)} samples saved to {shard_path}"
    
    # Prepare shard information
    shard_info = []
    for i in range(num_shards):
        start = i * shard_size
        end = min(len(confirmed_dataset), (i + 1) * shard_size)
        shard_info.append((i, start, end))
    
    # Write shards in parallel
    print(f"Writing {num_shards} shards in parallel...")
    with ProcessPoolExecutor(max_workers=NUM_CORES) as executor:
        results = list(tqdm(
            executor.map(write_shard, shard_info),
            total=len(shard_info),
            desc="Writing shards"
        ))
    
    # Print results
    for result in results:
        if result:
            print(f"🔹 {result}")
    
    print(f"\n✅ All confirmed data saved in {num_shards} shards in `{output_dir}/`")
    
    return num_shards

async def upload_to_hf(num_shards):
    """Upload to Hugging Face Hub"""
    print(f"\n🚀 Pushing dataset to Hugging Face Hub as '{HF_DATASET_NAME}'...")
    try:
        api = HfApi(token=os.getenv("HF_TOKEN"))
        
        # Create repository
        try:
            create_repo(
                repo_id=HF_DATASET_NAME,
                repo_type="dataset",
                private=HF_PRIVATE,
                exist_ok=True
            )
            print(f"✅ Repository '{HF_DATASET_NAME}' created/verified")
        except Exception as e:
            print(f"⚠️  Repository creation failed: {e}")
            return
        
        # Create dataset info
        dataset_info = {
            "dataset_name": HF_DATASET_NAME,
            "description": "Persian Common Voice confirmed samples for Whisper fine-tuning",
            "total_samples": len(confirmed),
            "num_shards": num_shards,
            "audio_format": "int16 PCM, 16kHz",
            "columns": ["audio", "text"],
            "source_dataset": "Ashegh-Sad-Warrior/Persian_Common_Voice_17_0",
            "processing": "Vosk API batch confirmation (optimized for 192 cores)"
        }
        
        info_path = os.path.join(output_dir, "dataset_info.json")
        with open(info_path, 'w', encoding='utf-8') as f:
            json.dump(dataset_info, f, indent=2, ensure_ascii=False)
        
        # Upload folder
        api.upload_folder(
            folder_path=output_dir,
            repo_id=HF_DATASET_NAME,
            repo_type="dataset",
        )
        
        print(f"🎉 Dataset successfully pushed to: https://huggingface.co/datasets/{HF_DATASET_NAME}")
        
    except Exception as e:
        print(f"❌ Failed to push to Hugging Face: {e}")

async def main():
    """Main function"""
    start_time = time.time()
    
    print(f"🚀 Starting optimized processing with {NUM_CORES} cores")
    print(f"📊 Dataset size: {len(ds)} samples")
    print(f"⚙️  Batch size: {BATCH_SIZE}")
    print(f"🔄 Max concurrent requests: {MAX_CONCURRENT_REQUESTS}")
    
    # Process dataset
    confirmed = await process_dataset_async()
    
    # Save data
    num_shards = save_confirmed_data_parallel(confirmed)
    
    # Upload to HF
    await upload_to_hf(num_shards)
    
    end_time = time.time()
    print(f"\n⏱️  Total processing time: {end_time - start_time:.2f} seconds")
    print(f"📈 Processing rate: {len(ds) / (end_time - start_time):.2f} samples/second")

if __name__ == "__main__":
    # Set multiprocessing start method for better performance
    mp.set_start_method('spawn', force=True)
    
    # Run the async main function
    asyncio.run(main())