from datasets import load_dataset, Audio, Dataset
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

# Load the dataset with audio decoding
print("Loading dataset...")
ds = load_dataset(
    "Ashegh-Sad-Warrior/Persian_Common_Voice_17_0",
    split="validated",
    streaming=False
).cast_column("audio", Audio(sampling_rate=16000))

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

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

# Hugging Face configuration
HF_DATASET_NAME = "dpr2000/persian-cv17-confirmed"  # Change this to your desired dataset name
HF_PRIVATE = True  # Set to True for private dataset

def save_flac(audio_array, path):
    sf.write(path, audio_array, 16000, format="FLAC")

print("Processing batches...")
for i in tqdm(range(0, len(ds), batch_size)):
    batch = ds[i:i+batch_size]
    files = {}
    references = []
    temp_flacs = []
    audio_arrays = []
    # Fix: batch is a dict of lists
    for j in range(len(batch["audio"])):
        audio = batch["audio"][j]
        flac_path = f"temp_{i+j}.flac"
        save_flac(audio["array"], flac_path)
        files[f"audio{j}"] = open(flac_path, "rb")
        references.append(batch["sentence"][j])
        temp_flacs.append(flac_path)
        audio_arrays.append(audio["array"])  # Store the array for confirmed
    data = {"references": json.dumps(references)}
    try:
        response = requests.post(API_URL, files=files, data=data, timeout=120)
        if response.status_code == 200:
            resp_json = response.json()
            if "results" in resp_json:
                results = resp_json["results"]
            else:
                print(f"Batch {i} failed: 'results' key missing in response: {resp_json}")
                results = [None] * len(references)
        else:
            print(f"Batch {i} failed: HTTP {response.status_code} - {response.text}")
            results = [None] * len(references)
    except Exception as e:
        print(f"Batch {i} failed: {e}")
        results = [None] * len(references)
    for j, result in enumerate(results):
        if result and result.get("confirmed"):
            # Save confirmed audio array and transcription
            confirmed.append({"audio": audio_arrays[j], "transcription": references[j]})
            os.remove(temp_flacs[j])
        else:
            os.remove(temp_flacs[j])
    for f in files.values():
        f.close()

# Save confirmed data using sharding approach
if confirmed:
    print(f"\n🔄 Saving {len(confirmed)} confirmed samples...")
    
    # Convert confirmed data to HuggingFace dataset format
    def extract_minimal(example):
        # Convert float32 audio (range -1.0 to 1.0) to int16 (range -32768 to 32767)
        audio_float32 = np.array(example["audio"], dtype=np.float32)
        # Ensure audio is in valid range and scale to int16
        audio_float32 = np.clip(audio_float32, -1.0, 1.0)
        audio_int16 = (audio_float32 * 32767).astype(np.int16)
        return {
            "audio": audio_int16.tobytes(),  # Store as int16 bytes, compatible with Whisper
            "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)
    
    # Sharding parameters
    num_shards = min(25, len(confirmed))  # Don't create more shards than samples
    shard_size = len(confirmed_dataset) // num_shards + 1
    
    # Write each shard separately
    for i in range(num_shards):
        start = i * shard_size
        end = min(len(confirmed_dataset), (i + 1) * shard_size)
        
        if start >= len(confirmed_dataset):
            break
            
        shard = confirmed_dataset.select(range(start, end))
        table = pa.Table.from_pandas(shard.to_pandas())  # Convert to PyArrow table
        
        shard_path = os.path.join(output_dir, f"confirmed_shard_{i:02}.parquet")
        
        pq.write_table(
            table,
            shard_path,
            compression="zstd",
            compression_level=22,     # Maximum compression
            use_dictionary=True,
            version="2.6"
        )
        
        print(f"🔹 Shard {i+1}/{num_shards}: {len(shard)} samples saved")
    
    print(f"\n✅ All confirmed data saved in {num_shards} shards in `{output_dir}/`")
    
    # Push to Hugging Face Hub
    print(f"\n🚀 Pushing dataset to Hugging Face Hub as '{HF_DATASET_NAME}'...")
    try:
        # Initialize HF API with token
        api = HfApi(token=os.getenv("HF_TOKEN"))
        
        # Create the repository (private if specified)
        repo_created = False
        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")
            repo_created = True
        except Exception as e:
            print(f"⚠️  Repository creation failed: {e}")
            print("💡 Please create the repository manually on Hugging Face Hub first")
            print(f"💡 Or change HF_DATASET_NAME to use your own username")
            print("💡 Skipping upload due to repository creation failure")
        
        if not repo_created:
            print("💡 Skipping upload due to repository creation failure")
        else:
            # Create dataset info file
            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"
            }
            
            # Save dataset info to output directory
            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 entire folder using 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}")
        print("💡 Make sure you have HF_TOKEN environment variable set")
        print("💡 Set it with: export HF_TOKEN=your_token_here")

else:
    print("❌ No confirmed samples to save")