import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import pandas as pd
from pathlib import Path
import requests
import zipfile
import os

class DatasetRegistry:
    """Registry for managing dataset URLs and metadata"""
    
    def __init__(self):
        self.datasets = {
            "cifar10": "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz",
            "imagenet_sample": "https://download.pytorch.org/tutorial/data.zip",
            "custom_data": "https://example.com/datasets/custom.zip"
        }
    
    def get_url(self, dataset_name: str) -> str:
        """Get dataset URL by name"""
        if dataset_name not in self.datasets:
            raise ValueError(f"Dataset '{dataset_name}' not found in registry")
        return self.datasets[dataset_name]
    
    def download_dataset(self, dataset_name: str, data_dir: str = "data"):
        """Download and extract dataset"""
        url = self.get_url(dataset_name)
        data_path = Path(data_dir)
        data_path.mkdir(exist_ok=True)
        
        print(f"Downloading {dataset_name} from {url}...")
        response = requests.get(url, stream=True)
        
        # Save the file
        filename = url.split('/')[-1]
        filepath = data_path / filename
        with open(filepath, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        
        # Extract if it's a zip file
        if filename.endswith('.zip'):
            with zipfile.ZipFile(filepath, 'r') as zip_ref:
                zip_ref.extractall(data_path)
        
        print(f"Dataset {dataset_name} downloaded and extracted to {data_path}")
        return data_path

class StandardDataset(Dataset):
    """Standard PyTorch Dataset wrapper"""
    
    def __init__(self, data_path: str, transform=None):
        self.data_path = Path(data_path)
        self.transform = transform
        self.data = self._load_data()
    
    def _load_data(self):
        # Override this method in subclasses
        raise NotImplementedError
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        sample = self.data[idx]
        if self.transform:
            sample = self.transform(sample)
        return sample

class CIFAR10Dataset(StandardDataset):
    """CIFAR-10 dataset implementation"""
    
    def _load_data(self):
        # Standard CIFAR-10 loading logic
        import pickle
        
        data = []
        for batch_file in self.data_path.glob("cifar-10-batches-py/data_batch_*"):
            with open(batch_file, 'rb') as f:
                batch = pickle.load(f, encoding='bytes')
                data.extend(list(zip(batch[b'data'], batch[b'labels'])))
        
        return data
    
    def __getitem__(self, idx):
        img_data, label = self.data[idx]
        img = img_data.reshape(3, 32, 32).transpose(1, 2, 0)  # HWC format
        
        if self.transform:
            img = self.transform(img)
        
        return img, label

def get_dataloader(dataset_name: str, batch_size: int = 32, transform=None):
    """Get a DataLoader for a registered dataset"""
    
    # Initialize registry and download dataset
    registry = DatasetRegistry()
    data_path = registry.download_dataset(dataset_name)
    
    # Create appropriate dataset
    if dataset_name == "cifar10":
        dataset = CIFAR10Dataset(data_path, transform=transform)
    else:
        # Generic dataset for other types
        dataset = StandardDataset(data_path, transform=transform)
    
    # Create and return DataLoader
    return DataLoader(dataset, batch_size=batch_size, shuffle=True)

if __name__ == "__main__":
    # Example usage
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    
    dataloader = get_dataloader("cifar10", batch_size=64, transform=transform)
    
    print(f"Dataset loaded with {len(dataloader)} batches")
    
    # Test loading a batch
    for images, labels in dataloader:
        print(f"Batch shape: {images.shape}, Labels: {labels.shape}")
        break
