"""
Dataset Splits API
"""

from typing import Optional, Dict, Any, List, Union
from pathlib import Path
from .base import BaseAPI
import pandas as pd


class DatasetSplitsAPI(BaseAPI):
    """Dataset split management endpoints"""
    
    def list(
        self,
        dataset_version_id: Optional[int] = None,
        config_name: Optional[str] = None,
        name: Optional[str] = None,
        page: int = 1,
        page_size: int = 20,
    ) -> Dict[str, Any]:
        """List dataset splits. Filter: dataset_version_id, config_name, name."""
        params = {"page": page, "page_size": min(page_size, 100)}
        if dataset_version_id is not None:
            params["dataset_version_id"] = dataset_version_id
        if config_name:
            params["config_name"] = config_name
        if name:
            params["name"] = name
        return self.client._request("GET", "dataset_splits/", params=params)
    
    def get(self, split_id: int) -> Dict[str, Any]:
        """Get dataset split by ID (GET /dataset_splits/by-id/{split_id})."""
        return self.client._request("GET", f"dataset_splits/by-id/{split_id}")
    
    def create(
        self,
        dataset_version_id: int,
        name: str,
        config_name: str,
    ) -> Dict[str, Any]:
        """
        Create a new dataset split (metadata only).
        Backend chỉ lưu bản ghi split (dataset_version_id, name, config_name).
        Để gán file vào split, dùng create_split_from_list() sau khi có danh sách file
        (vd: lấy từ get_mfcc_files() rồi chia bằng pandas/sklearn).

        Args:
            dataset_version_id: Dataset version ID
            name: Split name (e.g. "train", "val", "test")
            config_name: Config name (e.g. "config_70_15_15")

        Returns:
            Created split with dataset_splits_id, dataset_version_id, name, config_name, created_at, is_generated
        """
        data = {
            "dataset_version_id": dataset_version_id,
            "name": name,
            "config_name": config_name,
        }
        return self.client._request("POST", "dataset_splits/", json=data)
    
    def update(
        self,
        split_id: int,
        name: Optional[str] = None,
        config_name: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Update dataset split (name, config_name). Không đổi config_name nếu split đã generate."""
        data = {}
        if name is not None:
            data["name"] = name
        if config_name is not None:
            data["config_name"] = config_name
        return self.client._request("PUT", f"dataset_splits/{split_id}", json=data)
    
    def delete(self, split_id: int) -> None:
        """Delete dataset split"""
        self.client._request("DELETE", f"dataset_splits/{split_id}")
    
    def download(self, split_id: int, output_path: str) -> str:
        """
        Download dataset split as ZIP file.
        
        Args:
            split_id: Dataset split ID
            output_path: Path to save the ZIP file
            
        Returns:
            Path to downloaded file
        """
        response = self.client._request(
            "GET",
            f"dataset_splits/{split_id}/download",
            stream=True
        )
        
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        with open(output_path, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        
        return str(output_path)
    
    def list_files(
        self,
        split_id: int,
        file_type: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        List all files (.wav or .npz) in dataset split.
        
        Args:
            split_id: Dataset split ID
            file_type: Filter by file type ('wav' or 'npz'), None for all
            
        Returns:
            Dictionary with file list and metadata
        """
        params = {}
        if file_type:
            params["file_type"] = file_type
        return self.client._request(
            "GET",
            f"dataset_splits/{split_id}/files",
            params=params
        )
    
    def download_file(
        self,
        split_id: int,
        file_id: int,
        file_type: str,
        output_path: Optional[str] = None,
    ) -> str:
        """
        Download a specific file (.wav or .npz) from dataset split via API.
        Uses API endpoint only - no direct MinIO access.
        Backend streams file directly, no redirects to MinIO.
        
        Args:
            split_id: Dataset split ID
            file_id: File ID
            file_type: File type ('wav' or 'npz')
            output_path: Path to save file (required)
            
        Returns:
            Path to downloaded file
        """
        if not output_path:
            raise ValueError("output_path is required for download_file")
        
        params = {"file_type": file_type}
        response = self.client._request(
            "GET",
            f"dataset_splits/{split_id}/files/{file_id}/download",
            params=params,
            stream=True
        )
        
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        with open(output_path, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
        
        return str(output_path)
    
    def download_all_files(
        self,
        split_id: int,
        file_type: str,
        format: str = "zip",
        output_path: Optional[str] = None,
    ) -> Union[str, Dict[str, Any]]:
        """
        Download all files (.wav or .npz) from dataset split.
        
        Args:
            split_id: Dataset split ID
            file_type: File type ('wav' or 'npz')
            format: Download format ('zip' or 'urls')
            output_path: Path to save ZIP file (required if format='zip')
            
        Returns:
            Path to downloaded ZIP file or dictionary with URLs
        """
        params = {"file_type": file_type, "format": format}
        
        if format == "zip":
            if not output_path:
                raise ValueError("output_path is required when format='zip'")
            
            response = self.client._request(
                "GET",
                f"dataset_splits/{split_id}/files/download-all",
                params=params,
                stream=True
            )
            
            output_path = Path(output_path)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            
            with open(output_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
            
            return str(output_path)
        else:
            # Return URLs list
            return self.client._request(
                "GET",
                f"dataset_splits/{split_id}/files/download-all",
                params=params
            )
    
    def generate(self, split_id: int) -> Dict[str, Any]:
        """Generate dataset split (triggers background job)"""
        return self.client._request(
            "POST",
            f"dataset_splits/{split_id}/generate"
        )
    
    def get_mfcc_files(
        self,
        dataset_version_id: int,
        feature_type_id: Optional[int] = None,
        keyword_ids: Optional[List[int]] = None,
    ) -> pd.DataFrame:
        """
        Get list of all MFCC feature files with labels as pandas DataFrame.
        
        Args:
            dataset_version_id: Dataset version ID
            feature_type_id: Optional filter by feature type ID
            keyword_ids: Optional filter by keyword IDs
            
        Returns:
            pandas DataFrame with columns: feature_id, file_name, minio_path, 
            keyword_id, keyword_name, sample_id, feature_type_id
        """
        params = {"dataset_version_id": dataset_version_id}
        if feature_type_id:
            params["feature_type_id"] = feature_type_id
        if keyword_ids:
            params["keyword_ids"] = keyword_ids
        
        response = self.client._request("GET", "dataset_splits/mfcc-files", params=params)
        
        # Convert to DataFrame
        files = response.get("files", [])
        if not files:
            return pd.DataFrame(columns=[
                "feature_id", "file_name", "minio_path", 
                "keyword_id", "keyword_name", "sample_id", "feature_type_id"
            ])
        
        return pd.DataFrame(files)
    
    def create_split_from_list(
        self,
        dataset_version_id: int,
        split_name: str,
        config_name: str,
        files: List[Dict[str, Any]],
        create_merged_keywords: Optional[Dict[str, List[int]]] = None,
    ) -> Dict[str, Any]:
        """
        Create dataset split from a list of file assignments.
        
        Args:
            dataset_version_id: Dataset version ID
            split_name: Split name (e.g., "train", "val", "test")
            config_name: Configuration name (e.g., "config_70_15_15")
            files: List of file assignments, each with:
                - feature_id: int (required)
                - keyword_id: int (required)
                - keyword_name: str (optional, for new merged keywords)
            create_merged_keywords: Optional dict mapping new keyword names to 
                list of source keyword IDs to merge (e.g., {"unknown": [1, 2, 3]})
        
        Returns:
            Created dataset split response
        """
        payload = {
            "dataset_version_id": dataset_version_id,
            "split_name": split_name,
            "config_name": config_name,
            "files": files,
        }
        
        if create_merged_keywords:
            payload["create_merged_keywords"] = create_merged_keywords
        
        return self.client._request(
            "POST",
            "dataset_splits/create-from-list",
            json=payload
        )
