"""
Data loading and preparation utilities
"""

import os
import zipfile
import numpy as np
from pathlib import Path
from typing import Optional, Tuple, Dict, Any, List
from kwslib import DatasetSplitFilesClient
from kwslib.utils.load_npz_files import load_npz_files


class DataLoader:
    """Helper class to load and prepare training data"""
    
    def __init__(self, api_client, files_client: DatasetSplitFilesClient):
        self.api = api_client
        self.files_client = files_client
    
    def download_and_load_data(
        self,
        train_split_id: int,
        test_split_id: int,
        base_dir: str = "data"
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, Dict[str, int]]:
        """
        Download and load training and test data.
        
        Args:
            train_split_id: Train split ID
            test_split_id: Test split ID
            base_dir: Base directory for storing downloaded files
        
        Returns:
            Tuple of (X_train, y_train, X_test, y_test, label_to_id)
        """
        # Get file information
        print("\n=== Getting file information ===")
        train_files_info = self.files_client.list_files(split_id=train_split_id, file_type="npz")
        test_files_info = self.files_client.list_files(split_id=test_split_id, file_type="npz")
        
        train_files = train_files_info.get('files', [])
        test_files = test_files_info.get('files', [])
        
        # Download training data
        print(f"\n=== Downloading training data (split {train_split_id}) ===")
        train_dir = os.path.join(base_dir, f"train_split_{train_split_id}")
        os.makedirs(train_dir, exist_ok=True)
        self.files_client.download_all_npz(split_id=train_split_id, output_dir=train_dir)
        
        # Download test data
        print(f"\n=== Downloading test data (split {test_split_id}) ===")
        test_dir = os.path.join(base_dir, f"test_split_{test_split_id}")
        os.makedirs(test_dir, exist_ok=True)
        self.files_client.download_all_npz(split_id=test_split_id, output_dir=test_dir)
        
        # Load training data (creates label mapping)
        print("\n=== Loading training data ===")
        print("Using derivative_label from split for training labels")
        X_train, y_train, label_to_id = load_npz_files(
            train_dir,
            api_client=self.api,
            file_info_map=train_files
        )
        
        # Load test data (reuses label mapping from train)
        print("\n=== Loading test data ===")
        print("Using derivative_label from split for test labels")
        print("Reusing label mapping from training data")
        X_test, y_test, label_to_id = load_npz_files(
            test_dir,
            api_client=self.api,
            file_info_map=test_files,
            label_to_id=label_to_id  # Reuse mapping from train
        )
        
        return X_train, y_train, X_test, y_test, label_to_id
    
    def download_and_load_data_zip(
        self,
        train_split_id: int,
        test_split_id: int,
        base_dir: str = "data",
        keep_zip: bool = False,
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, Dict[str, int]]:
        """
        Download and load training and test data using ZIP endpoints for better performance.

        Thay vì tải từng file .npz một, method này:
        - Gọi backend để tải toàn bộ file của split dưới dạng 1 file .zip
        - Giải nén ra thư mục local
        - Gọi lại load_npz_files giống như download_and_load_data
        """
        # Get file information (để map label / metadata)
        print("\n=== Getting file information ===")
        train_files_info = self.files_client.list_files(split_id=train_split_id, file_type="npz")
        test_files_info = self.files_client.list_files(split_id=test_split_id, file_type="npz")

        train_files = train_files_info.get("files", [])
        test_files = test_files_info.get("files", [])

        # Paths
        base_dir = Path(base_dir)
        base_dir.mkdir(parents=True, exist_ok=True)

        train_dir = base_dir / f"train_split_{train_split_id}"
        test_dir = base_dir / f"test_split_{test_split_id}"
        train_dir.mkdir(parents=True, exist_ok=True)
        test_dir.mkdir(parents=True, exist_ok=True)

        train_zip = base_dir / f"train_split_{train_split_id}.zip"
        test_zip = base_dir / f"test_split_{test_split_id}.zip"

        # Download ZIPs (1 request / split)
        print(f"\n=== Downloading training data as ZIP (split {train_split_id}) ===")
        self.files_client.download_all_files_zip(
            split_id=train_split_id,
            file_type="npz",
            output_path=str(train_zip),
        )

        print(f"\n=== Downloading test data as ZIP (split {test_split_id}) ===")
        self.files_client.download_all_files_zip(
            split_id=test_split_id,
            file_type="npz",
            output_path=str(test_zip),
        )

        # Extract ZIPs
        print(f"\n=== Extracting training ZIP to {train_dir} ===")
        with zipfile.ZipFile(train_zip, "r") as zf:
            zf.extractall(train_dir)

        print(f"\n=== Extracting test ZIP to {test_dir} ===")
        with zipfile.ZipFile(test_zip, "r") as zf:
            zf.extractall(test_dir)

        if not keep_zip:
            try:
                train_zip.unlink(missing_ok=True)  # type: ignore[call-arg]
                test_zip.unlink(missing_ok=True)   # type: ignore[call-arg]
            except TypeError:
                # Python <3.8 không có missing_ok
                if train_zip.exists():
                    train_zip.unlink()
                if test_zip.exists():
                    test_zip.unlink()

        # Load training data (creates label mapping)
        print("\n=== Loading training data ===")
        print("Using derivative_label from split for training labels")
        X_train, y_train, label_to_id = load_npz_files(
            str(train_dir),
            api_client=self.api,
            file_info_map=train_files,
        )

        # Load test data (reuses label mapping from train)
        print("\n=== Loading test data ===")
        print("Using derivative_label from split for test labels")
        print("Reusing label mapping from training data")
        X_test, y_test, label_to_id = load_npz_files(
            str(test_dir),
            api_client=self.api,
            file_info_map=test_files,
            label_to_id=label_to_id,
        )

        return X_train, y_train, X_test, y_test, label_to_id
    
    def prepare_for_cnn(
        self,
        X_train: np.ndarray,
        X_test: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Prepare data for CNN training (reshape and normalize).
        
        Args:
            X_train: Training features
            X_test: Test features
        
        Returns:
            Tuple of (normalized_X_train, normalized_X_test)
        """
        # Reshape if needed - CNN expects (samples, time_steps, features)
        if len(X_train.shape) == 2:
            # If 2D, assume it's (samples, features) - add time dimension
            # Reshape to (samples, time_steps=1, features)
            X_train = np.expand_dims(X_train, axis=1)
            X_test = np.expand_dims(X_test, axis=1)
        elif len(X_train.shape) == 3:
            # Already 3D (samples, time, frequency)
            pass
        else:
            raise ValueError(f"Unexpected feature shape: {X_train.shape}")
        
        print(f"\nFinal feature shape: {X_train.shape}")
        
        # Normalize features
        mean = X_train.mean(axis=(0, 1), keepdims=True)
        std = X_train.std(axis=(0, 1), keepdims=True) + 1e-8
        X_train = (X_train - mean) / std
        X_test = (X_test - mean) / std
        
        return X_train, X_test


def load_train_val_test_mfcc(
    api_client,
    files_client: DatasetSplitFilesClient,
    train_split_id: int,
    val_split_id: Optional[int] = None,
    test_split_id: Optional[int] = None,
    base_dir: str = "data",
    n_mels: Optional[int] = None,
) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray], Dict[str, int]]:
    """
    Wrapper high-level để tải feature MFCC/mel cho train/val/test và chuẩn hóa shape.

    Args:
        api_client: KWSClient đã login.
        files_client: DatasetSplitFilesClient.
        train_split_id: ID split train.
        val_split_id: ID split val (nếu None sẽ dùng test_split_id như val).
        test_split_id: ID split test (tùy chọn).
        base_dir: Thư mục gốc để lưu ZIP/npz.
        n_mels: Nếu không None, chuẩn hóa shape về (samples, n_mels, time).

    Returns:
        (X_train, y_train, X_val, y_val, X_test, y_test, label_to_id)
    """
    loader = DataLoader(api_client, files_client)

    # Train + Val: dùng sẵn download_and_load_data_zip cho train/val
    if val_split_id is None:
        raise ValueError("val_split_id phải được cung cấp để load train/val.")

    X_train, y_train, X_val, y_val, label_to_id = loader.download_and_load_data_zip(
        train_split_id=train_split_id,
        test_split_id=val_split_id,
        base_dir=base_dir,
    )

    # Chuẩn hóa shape (samples, n_mels, time) nếu cần
    def _normalize_shape(X: np.ndarray) -> np.ndarray:
        if X is None or n_mels is None:
            return X
        if X.ndim == 2:
            X = np.expand_dims(X, axis=1)
        if X.ndim == 3 and X.shape[1] != n_mels and X.shape[2] == n_mels:
            X = np.transpose(X, (0, 2, 1))
        return X

    X_train = _normalize_shape(X_train)
    X_val = _normalize_shape(X_val)

    X_test: Optional[np.ndarray] = None
    y_test: Optional[np.ndarray] = None

    # Nếu có test_split_id, tải thêm test
    if test_split_id is not None:
        base = Path(base_dir)
        base.mkdir(parents=True, exist_ok=True)
        test_dir = base / f"test_split_{test_split_id}"
        test_dir.mkdir(parents=True, exist_ok=True)
        test_zip = base / f"test_split_{test_split_id}.zip"

        print(f"\n=== Getting file information (test split {test_split_id}) ===")
        test_files_info = files_client.list_files(split_id=test_split_id, file_type="npz")
        test_files = test_files_info.get("files", [])

        print(f"\n=== Downloading test data as ZIP (split {test_split_id}) ===")
        files_client.download_all_files_zip(
            split_id=test_split_id,
            file_type="npz",
            output_path=str(test_zip),
        )

        print(f"\n=== Extracting test ZIP to {test_dir} ===")
        with zipfile.ZipFile(test_zip, "r") as zf:
            zf.extractall(test_dir)
        try:
            test_zip.unlink()
        except FileNotFoundError:
            pass

        print("\n=== Loading test data ===")
        print("Using derivative_label from split for test labels")
        print("Reusing label mapping from training data")
        X_test, y_test, _ = load_npz_files(
            str(test_dir),
            api_client=api_client,
            file_info_map=test_files,
            label_to_id=label_to_id,
        )
        X_test = _normalize_shape(X_test)

    return X_train, y_train, X_val, y_val, X_test, y_test, label_to_id

