from torch.utils.data import Dataset, DataLoader
from typing import *
from dataclasses import dataclass, field
from PIL import Image
from utils import parse_structure

import os
import lightning.pytorch as pl
import numpy as np
import torch


class BaseDataset(Dataset):
    def __init__(self, root_dir: str, image_size: Tuple[int, int]) -> None:
        self.root_dir = root_dir
        self.image_size = image_size
        self.classes = {folder: idx for idx, folder in enumerate(os.listdir(root_dir))}
        self.image_paths = []
        self.labels = []

        for class_name, class_idx in self.classes.items():
            class_dir = os.path.join(root_dir, class_name)
            for img_name in os.listdir(class_dir):
                img_path = os.path.join(class_dir, img_name)
                self.image_paths.append(img_path)
                self.labels.append(class_idx)

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
        img_path = self.image_paths[idx]
        label = self.labels[idx]
        image = Image.open(img_path).convert("RGB")
        image = image.resize(self.image_size)
        image = np.array(image)
        image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
        return image, label


@dataclass
class BaseDatasetConfig:
    data_source: str = ''
    train_path:str = ''
    valid_path:str = ''
    test_path:str = ''
    batch_size:int = 32
    shuffle:bool = True
    num_workers:int = 24
    image_size:Tuple[int, int] = (224, 224)

class BaseDataModule(pl.LightningDataModule):
    cfg: BaseDatasetConfig

    def __init__(self, cfg: BaseDatasetConfig) -> None:
        super().__init__()
        self.cfg:BaseDatasetConfig = parse_structure(BaseDatasetConfig, cfg)
        self.train_path = cfg.train_path
        self.valid_path = cfg.valid_path
        self.test_path = cfg.test_path
        self.img_size = cfg.image_size

    def setup(self, stage=None) -> None:
        if stage in [None, "fit"]:
            self.train_dataset = BaseDataset(self.train_path, self.img_size)
        if stage in [None, "fit", "validate"]:
            self.val_dataset = BaseDataset(self.valid_path, self.img_size)
        if stage in [None, "test", "predict"]:
            self.test_dataset = BaseDataset(self.test_path, self.img_size)

    def general_loader(self, dataset, batch_size) -> DataLoader:
        return DataLoader(
            dataset, 
            num_workers=self.cfg.num_workers, 
            batch_size=batch_size
        )

    def train_dataloader(self) -> DataLoader:
        return DataLoader(
            self.train_dataset, 
            num_workers=self.cfg.num_workers, 
            batch_size=self.cfg.batch_size, 
            shuffle=self.cfg.shuffle
        )

    def val_dataloader(self) -> DataLoader:
        return DataLoader(
            self.val_dataset, 
            num_workers=self.cfg.num_workers, 
            batch_size=self.cfg.batch_size
        )

    def test_dataloader(self) -> DataLoader:
        return DataLoader(
            self.test_dataset, 
            num_workers=self.cfg.num_workers, 
            batch_size=self.cfg.batch_size
        )