import os
import logging
from typing import List, Tuple, Optional, Union, Dict, Any, Callable
from pathlib import Path

import torch
from torch.utils.data import Dataset
from PIL import Image
import numpy as np

try:
    import yaml
except ImportError:
    yaml = None

logger = logging.getLogger(__name__)


class SegDataset(Dataset):
    """
    PyTorch Dataset for segmentation tasks.

    This dataset supports multiple configuration formats:
    - **Folder structure**: expects `images/` and `masks/` subdirectories.
    - **Text file**: each line contains `image_path mask_path`.
    - **YAML file**: structured configuration with optional class names.
    - **Dictionary**: same as YAML but passed directly.

    The dataset loads images as RGB tensors normalized to [0,1] (CxHxW) and masks
    as long tensors (HxW) with class indices starting from 0.

    Args:
        config (str or dict): Path to dataset configuration (folder, .txt, .yaml) or a dict with keys.
        format (str, optional): Explicit format: 'folder', 'txt', 'yaml'. If None, inferred from extension.
        split (str, optional): For folder format, specify 'train' or 'val' to use split subfolders or split files.
        transforms (callable, optional): Transform(s) to apply to image and mask.
            Should accept and return (image, mask) tuple.
        image_suffix (str, optional): File extension for images (e.g., '.png', '.jpg'). Used in folder mode.
        mask_suffix (str, optional): File extension for masks. If None, uses same as image.
        validate (bool): If True, check that all files exist. Default True.
        **kwargs: Additional arguments passed to the YAML parser.

    Attributes:
        image_paths (list of str): List of image file paths.
        mask_paths (list of str): Corresponding mask file paths.
        transforms (callable or None): Transform function applied to each sample.
        class_names (list of str): Optional class names (loaded from YAML/dict).
        num_classes (int): Number of classes (inferred from masks or config).

    Examples:
        **Folder format with split subdirectories:**
        >>> dataset = SegDataset(
        ...     config='path/to/data',
        ...     format='folder',
        ...     split='train',
        ...     image_suffix='.jpg',
        ...     mask_suffix='.png'
        ... )

        **Folder format with split file:**
        >>> # Requires data/train.txt with filenames (one per line)
        >>> dataset = SegDataset(
        ...     config='path/to/data',
        ...     split='train',
        ...     image_suffix='.jpg'
        ... )

        **Text file format:**
        >>> dataset = SegDataset(config='path/to/files.txt', format='txt')

        **YAML configuration:**
        >>> # dataset.yaml content:
        >>> # images: /path/to/images
        >>> # masks: /path/to/masks
        >>> # class_names: ['background', 'cat', 'dog']
        >>> dataset = SegDataset(config='dataset.yaml')

        **Dictionary configuration:**
        >>> config_dict = {
        ...     'image_list': ['img1.png', 'img2.png'],
        ...     'mask_list': ['mask1.png', 'mask2.png'],
        ...     'class_names': ['bg', 'obj']
        ... }
        >>> dataset = SegDataset(config=config_dict, format='dict')

        **With transforms:**
        >>> from torchvision import transforms as T
        >>> transform = T.Compose([
        ...     T.ToTensor(),
        ...     T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ... ])
        >>> dataset = SegDataset(config='data', transforms=transform)

    Notes:
        - Masks are assumed to be single-channel images where pixel values are class indices.
        - The number of classes is automatically inferred from the maximum value in the first mask
          if not provided in the configuration.
    """

    def __init__(
        self,
        config: Union[str, dict],
        format: Optional[str] = None,
        split: Optional[str] = None,
        transforms: Optional[Callable] = None,
        image_suffix: str = '.png',
        mask_suffix: Optional[str] = None,
        validate: bool = True,
        **kwargs
    ):
        
        self.num_classes = None

        # Resolve format
        if format is None:
            if isinstance(config, str):
                ext = Path(config).suffix.lower()
                if ext == '.txt':
                    format = 'txt'
                elif ext in ('.yaml', '.yml'):
                    format = 'yaml'
                else:
                    format = 'folder' 
            else:
                format = 'dict'  

        if format == 'folder':
            if not isinstance(config, str):
                raise TypeError(f"For format='folder', config must be a string path, got {type(config).__name__}")
            image_paths, mask_paths = self._parse_folder(config, split, image_suffix, mask_suffix)
        elif format == 'txt':
            if not isinstance(config, str):
                raise TypeError(f"For format='txt', config must be a string path to a text file, got {type(config).__name__}")
            image_paths, mask_paths = self._parse_txt(config)
        elif format == 'yaml':
            if not isinstance(config, str):
                raise TypeError(f"For format='yaml', config must be a string path to a YAML file, got {type(config).__name__}")
            if yaml is None:
                raise ImportError("PyYAML is required to parse YAML files. Install with 'pip install pyyaml'.")
            with open(config, 'r') as f:
                cfg = yaml.safe_load(f)
            image_paths, mask_paths, extra = self._parse_dict(cfg, **kwargs)
            self.class_names = extra.get('class_names', [])
            self.num_classes = extra.get('num_classes', None)
        elif format == 'dict':
            if not isinstance(config, dict):
                raise TypeError(f"For format='dict', config must be a dictionary, got {type(config).__name__}")
            image_paths, mask_paths, extra = self._parse_dict(config, **kwargs)
            self.class_names = extra.get('class_names', [])
            self.num_classes = extra.get('num_classes', None)
        else:
            raise ValueError(f"Unsupported format: {format}")

        if validate:
            self._validate_paths(image_paths, mask_paths)

        self.image_paths = image_paths
        self.mask_paths = mask_paths
        self.transforms = transforms

        if self.num_classes is None and len(mask_paths) > 0:
            sample_mask = self._load_mask(mask_paths[0])
            self.num_classes = int(sample_mask.max().item()) + 1
            logger.info(f"Inferred number of classes from masks: {self.num_classes}")

    def __len__(self) -> int:
        """Return the total number of samples in the dataset."""
        return len(self.image_paths)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Retrieve the image and mask at the given index.

        Args:
            idx (int): Index of the sample.

        Returns:
            tuple: (image_tensor, mask_tensor) where image is float tensor [0,1] of shape (C,H,W)
                   and mask is long tensor of shape (H,W) with class indices.
        """
        image_path = self.image_paths[idx]
        mask_path = self.mask_paths[idx]

        image = self._load_image(image_path)
        mask = self._load_mask(mask_path)

        if self.transforms is not None:
            image, mask = self.transforms(image, mask)

        return image, mask

    @staticmethod
    def _parse_folder(
        root: str,
        split: Optional[str] = None,
        image_suffix: str = '.png',
        mask_suffix: Optional[str] = None
    ) -> Tuple[List[str], List[str]]:
        """
        Parse a folder structure.

        Two modes:
          - If split is None: expects `root/images/` and `root/masks/`.
          - If split is given: first checks for `root/split.txt` containing filenames (without extension),
            then looks in `root/split/images/` and `root/split/masks/`. If split.txt does not exist,
            assumes subdirectories `root/split/images/` and `root/split/masks/` and takes all files.

        Args:
            root (str): Root directory path.
            split (str, optional): Split name, e.g., 'train' or 'val'.
            image_suffix (str): Image file extension (including dot).
            mask_suffix (str, optional): Mask file extension; if None, uses image_suffix.

        Returns:
            tuple: (image_paths, mask_paths) as lists of strings.

        Raises:
            FileNotFoundError: If split file is specified but missing.
            NotADirectoryError: If required directories are missing.
        """
        root_path = Path(root)
        mask_suffix = mask_suffix or image_suffix

        if split is not None:
            split_file = root_path / f"{split}.txt"
            if split_file.exists():
                # Use split file listing filenames without extension
                with open(split_file, 'r') as f:
                    names = [line.strip() for line in f if line.strip()]
                images_dir = root_path / split / 'images'
                masks_dir = root_path / split / 'masks'
                image_paths = [str(images_dir / (name + image_suffix)) for name in names]
                mask_paths = [str(masks_dir / (name + mask_suffix)) for name in names]
            else:
                # Assume split subdirectories
                images_dir = root_path / split / 'images'
                masks_dir = root_path / split / 'masks'
                if not images_dir.is_dir():
                    raise NotADirectoryError(f"Images directory not found: {images_dir}")
                if not masks_dir.is_dir():
                    raise NotADirectoryError(f"Masks directory not found: {masks_dir}")
                image_files = sorted(images_dir.glob(f'*{image_suffix}'))
                mask_files = [masks_dir / (img_file.stem + mask_suffix) for img_file in image_files]
                image_paths = [str(p) for p in image_files]
                mask_paths = [str(p) for p in mask_files]
        else:
            images_dir = root_path / 'images'
            masks_dir = root_path / 'masks'
            if not images_dir.is_dir():
                raise NotADirectoryError(f"Images directory not found: {images_dir}")
            if not masks_dir.is_dir():
                raise NotADirectoryError(f"Masks directory not found: {masks_dir}")

            image_files = sorted([f for f in images_dir.iterdir() if f.suffix.lower() == image_suffix.lower()])
            mask_files = [masks_dir / (img_file.stem + mask_suffix) for img_file in image_files]

            image_paths = [str(p) for p in image_files]
            mask_paths = [str(p) for p in mask_files]

        return image_paths, mask_paths

    @staticmethod
    def _parse_txt(txt_path: str) -> Tuple[List[str], List[str]]:
        """
        Parse a text file with each line: image_path mask_path.

        Lines starting with '#' are ignored.

        Args:
            txt_path (str): Path to the text file.

        Returns:
            tuple: (image_paths, mask_paths) as lists of strings.

        Raises:
            ValueError: If a line does not contain exactly two paths.
        """
        image_paths = []
        mask_paths = []
        with open(txt_path, 'r') as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                parts = line.split()
                if len(parts) != 2:
                    raise ValueError(f"Line {line_num} in {txt_path} does not contain exactly two paths: {line}")
                img_p, mask_p = parts
                image_paths.append(img_p)
                mask_paths.append(mask_p)
        return image_paths, mask_paths

    @staticmethod
    def _parse_dict(cfg: dict, **kwargs) -> Tuple[List[str], List[str], Dict[str, Any]]:
        """
        Parse a dictionary configuration.

        Expected keys (one of the following combinations):
            - 'images' and 'masks': paths to directories (optionally with 'split' key).
            - 'image_list' and 'mask_list': explicit lists of paths.

        Optional keys:
            - 'class_names': list of class names.
            - 'num_classes': number of classes (if not inferred).

        Args:
            cfg (dict): Configuration dictionary.
            **kwargs: Additional arguments (unused, for compatibility).

        Returns:
            tuple: (image_paths, mask_paths, extra_dict) where extra_dict contains
                   'class_names' and 'num_classes'.

        Raises:
            ValueError: If required keys are missing.
            NotADirectoryError: If specified directories are not found.
        """
        image_paths = []
        mask_paths = []
        extra = {}

        if 'image_list' in cfg and 'mask_list' in cfg:
            image_paths = cfg['image_list']
            mask_paths = cfg['mask_list']
        elif 'images' in cfg and 'masks' in cfg:
            images_dir = Path(cfg['images'])
            masks_dir = Path(cfg['masks'])
            split = cfg.get('split', None)
            if split is not None:
                images_dir = images_dir / split
                masks_dir = masks_dir / split
            if not images_dir.is_dir():
                raise NotADirectoryError(f"Images directory not found: {images_dir}")
            if not masks_dir.is_dir():
                raise NotADirectoryError(f"Masks directory not found: {masks_dir}")

            image_files = sorted(images_dir.glob('*'))
            mask_files = []
            for img_file in image_files:
                # Try to find a mask with the same stem but any extension
                candidates = list(masks_dir.glob(img_file.stem + '.*'))
                if not candidates:
                    # Fallback: assume same filename
                    mask_files.append(masks_dir / img_file.name)
                else:
                    mask_files.append(candidates[0])
            image_paths = [str(p) for p in image_files]
            mask_paths = [str(p) for p in mask_files]
        else:
            raise ValueError("Dictionary must contain either ('images','masks') or ('image_list','mask_list')")

        extra['class_names'] = cfg.get('class_names', [])
        extra['num_classes'] = cfg.get('num_classes', None)

        return image_paths, mask_paths, extra

    @staticmethod
    def _load_image(path: str) -> torch.Tensor:
        """
        Load an image from disk and convert to a normalized float tensor.

        Args:
            path (str): Path to the image file.

        Returns:
            torch.Tensor: Image tensor of shape (C, H, W) with values in [0, 1].

        Raises:
            RuntimeError: If loading fails.
        """
        try:
            with Image.open(path) as img:
                img = img.convert('RGB')
                arr = np.array(img, dtype=np.float32) / 255.0
              
                tensor = torch.from_numpy(arr).permute(2, 0, 1)
                return tensor
        except Exception as e:
            raise RuntimeError(f"Failed to load image {path}: {e}")

    @staticmethod
    def _load_mask(path: str) -> torch.Tensor:
        """
        Load a segmentation mask from disk and convert to a long tensor.

        Args:
            path (str): Path to the mask file (grayscale PNG recommended).

        Returns:
            torch.Tensor: Mask tensor of shape (H, W) with class indices (int64).

        Raises:
            RuntimeError: If loading fails.
        """
        try:
            with Image.open(path) as msk:
                msk = msk.convert('L')  # ensure single channel
                arr = np.array(msk, dtype=np.int64)
                return torch.from_numpy(arr)
        except Exception as e:
            raise RuntimeError(f"Failed to load mask {path}: {e}")

    @staticmethod
    def _validate_paths(image_paths: List[str], mask_paths: List[str]):
        """
        Check that all image and mask files exist.

        Args:
            image_paths (list of str): List of image paths.
            mask_paths (list of str): List of mask paths.

        Raises:
            FileNotFoundError: If any file is missing.
        """
        missing = []
        for i, (img_p, mask_p) in enumerate(zip(image_paths, mask_paths)):
            if not os.path.exists(img_p):
                missing.append(f"Image {i}: {img_p}")
            if not os.path.exists(mask_p):
                missing.append(f"Mask {i}: {mask_p}")
        if missing:
            raise FileNotFoundError("Missing files:\n" + "\n".join(missing))
        logger.info(f"Validated {len(image_paths)} image-mask pairs.")

    @property
    def paths(self) -> List[Tuple[str, str]]:
        """Return list of (image_path, mask_path) tuples for all samples."""
        return list(zip(self.image_paths, self.mask_paths))

    def get_class_names(self) -> List[str]:
        """
        Return class names if available.

        Returns:
            list of str: Class names, or empty list if not set.
        """
        return getattr(self, 'class_names', [])

    def get_num_classes(self) -> int:
        """
        Return the number of classes.

        If not explicitly set, it is inferred from the maximum value in the first mask.

        Returns:
            int: Number of classes, or 0 if dataset is empty.
        """
        if hasattr(self, 'num_classes') and self.num_classes is not None:
            return self.num_classes
        if len(self.mask_paths) > 0:
            mask = self._load_mask(self.mask_paths[0])
            return int(mask.max().item()) + 1
        return 0