"""
Utilities for finding dataset splits
"""

import os
import json
from typing import Optional, Tuple, Dict, Any


class SplitFinder:
    """Helper class to find train and test split IDs"""
    
    def __init__(self, api_client):
        self.api = api_client
    
    def find_splits(
        self,
        train_split_id: Optional[int] = None,
        test_split_id: Optional[int] = None,
        config_name: Optional[str] = None,
        dataset_version_id: Optional[int] = None,
        split_ids_file: str = "data/split_ids.json"
    ) -> Tuple[int, int]:
        """
        Find train and test split IDs from various sources.
        
        Args:
            train_split_id: Explicit train split ID (highest priority)
            test_split_id: Explicit test split ID (highest priority)
            config_name: Config name to search for
            dataset_version_id: Dataset version ID for searching
            split_ids_file: Path to JSON file with saved split IDs
        
        Returns:
            Tuple of (train_split_id, test_split_id)
        
        Raises:
            ValueError: If splits cannot be found
        """
        # Try to load from split_ids.json file first
        if not train_split_id or not test_split_id:
            if os.path.exists(split_ids_file):
                try:
                    with open(split_ids_file, "r") as f:
                        split_data = json.load(f)
                        if config_name and split_data.get('config_name') == config_name:
                            saved_splits = split_data.get('splits', {})
                            if not train_split_id and 'train' in saved_splits:
                                train_split_id = saved_splits['train']
                                print(f"Loaded train split ID from {split_ids_file}: {train_split_id}")
                            if not test_split_id and 'test' in saved_splits:
                                test_split_id = saved_splits['test']
                                print(f"Loaded test split ID from {split_ids_file}: {test_split_id}")
                except Exception as e:
                    print(f"Warning: Could not load split IDs from {split_ids_file}: {e}")
        
        # If still missing, try to find from API
        if not train_split_id or not test_split_id:
            if config_name and dataset_version_id:
                print(f"\n=== Finding splits with config_name='{config_name}' ===")
                splits = self.api.dataset_splits.list(
                    dataset_version_id=dataset_version_id,
                    config_name=config_name,
                    page=1,
                    page_size=100
                )
                
                split_items = splits.get('items', []) if isinstance(splits, dict) else splits
                
                for split in split_items:
                    split_name = split.get('name', '').lower()
                    split_id = split.get('dataset_splits_id') or split.get('id')
                    
                    if split_name == 'train' and not train_split_id:
                        train_split_id = split_id
                        print(f"Found train split: ID={train_split_id}")
                    elif split_name == 'test' and not test_split_id:
                        test_split_id = split_id
                        print(f"Found test split: ID={test_split_id}")
                
                if not train_split_id:
                    print("Warning: Train split not found!")
                if not test_split_id:
                    print("Warning: Test split not found!")
            else:
                if not train_split_id or not test_split_id:
                    raise ValueError(
                        "Either provide --train-split-id and --test-split-id, "
                        "or --config-name and --dataset-version-id"
                    )
        
        if not train_split_id or not test_split_id:
            raise ValueError("Both train and test split IDs are required!")
        
        print(f"\nUsing splits:")
        print(f"  Train split ID: {train_split_id}")
        print(f"  Test split ID: {test_split_id}")
        
        return train_split_id, test_split_id
