"""
DICOM data models for handling DICOM files and series.
"""
import os
from typing import List, Dict, Any, Optional, Generator, Tuple
from concurrent.futures import ThreadPoolExecutor, as_completed
import pydicom
from pydicom.dataset import FileDataset
from pydicom.uid import generate_uid


class DICOMInstance:
    """
    Represents a single DICOM file instance.
    """
    def __init__(self, filepath: str):
        self.filepath = filepath
        self._dataset: Optional[FileDataset] = None

    @property
    def dataset(self) -> FileDataset:
        """
        Lazy-load the DICOM dataset.
        
        Returns:
            FileDataset: The pydicom dataset object
        """
        if self._dataset is None:
            try:
                self._dataset = pydicom.dcmread(self.filepath)
            except Exception as e:
                print(f"Error reading DICOM file {self.filepath}: {e}")
                raise
        return self._dataset

    def get_tag(self, tag: str) -> Any:
        """
        Get a DICOM tag value from the dataset.
        
        Args:
            tag: DICOM tag name
            
        Returns:
            Tag value or None if not found
        """
        try:
            return getattr(self.dataset, tag, None)
        except Exception as e:
            print(f"Error getting tag {tag} from {self.filepath}: {e}")
            return None

    @property
    def slice_location(self) -> Optional[float]:
        """获取 SliceLocation 标签值"""
        return self.get_tag("SliceLocation")

    @property
    def instance_number(self) -> Optional[int]:
        """获取 InstanceNumber 标签值"""
        return self.get_tag("InstanceNumber")

    def get_dicom_tag(self, tag: str, default: Any = None) -> Any:
        """
        获取 DICOM 标签值，支持默认值
        
        Args:
            tag: DICOM 标签名
            default: 默认值
            
        Returns:
            标签值或默认值
        """
        value = self.get_tag(tag)
        return value if value is not None else default

    def set_dicom_tag(self, tag: str, value: Any) -> None:
        """
        设置 DICOM 标签值
        
        Args:
            tag: DICOM 标签名
            value: 要设置的值
        """
        try:
            setattr(self.dataset, tag, value)
        except Exception as e:
            print(f"Error setting tag {tag} in {self.filepath}: {e}")

    @classmethod
    def copy_from_instance(cls, instance: "DICOMInstance") -> "DICOMInstance":
        """
        从现有实例创建副本
        
        Args:
            instance: 源实例
            
        Returns:
            新的 DICOMInstance 副本
        """
        new_instance = cls(instance.filepath)
        if instance._dataset is not None:
            # 创建数据集的副本
            new_instance._dataset = instance.dataset.copy()
        return new_instance


class Series:
    """
    Represents a DICOM series containing multiple instances.
    """
    def __init__(self, series_uid: str = None, folder_name: str = None, instances: List[DICOMInstance] = None):
        self.series_uid = series_uid
        self.folder_name = folder_name
        self.instances: List[DICOMInstance] = instances if instances is not None else []
        self._instance_number_dict: Optional[Dict[int, DICOMInstance]] = None

    def add_instance(self, instance: DICOMInstance) -> None:
        """
        Add a DICOM instance to this series.
        
        Args:
            instance: DICOMInstance to add
        """
        self.instances.append(instance)
        self._instance_number_dict = None  # 重置缓存

    @property
    def instance_number_dict(self) -> Dict[int, DICOMInstance]:
        """
        获取按 InstanceNumber 索引的实例字典
        
        Returns:
            字典，键为 InstanceNumber，值为 DICOMInstance
        """
        if self._instance_number_dict is None:
            self._instance_number_dict = {}
            for instance in self.instances:
                instance_num = instance.instance_number
                if instance_num is not None:
                    self._instance_number_dict[instance_num] = instance
        return self._instance_number_dict

    def get_dicom_tag(self, tag: str) -> Any:
        """
        Get a DICOM tag value from the first instance in the series.
        
        Args:
            tag: DICOM tag name
            
        Returns:
            Tag value or None if series is empty
        """
        if not self.instances:
            return None
        return self.instances[0].get_tag(tag)


def is_standard_dicom_file(file_path: str) -> bool:
    """
    快速检查文件是否为标准 DICOM 文件（只读取文件头，不完整读取文件）
    标准 DICOM 文件的前 128 字节是前导码，第 129-132 字节应该是 "DICM"
    
    Args:
        file_path: 文件路径
        
    Returns:
        bool: 是否为标准DICOM文件
    """
    try:
        with open(file_path, 'rb') as f:
            # 只读取前132字节进行快速检查
            header = f.read(132)
            
            if len(header) < 132:
                return False
                
            # 检查第128-131字节（索引从0开始）是否为"DICM"
            magic_word = header[128:132].decode('ascii', errors='ignore')
            return magic_word == 'DICM'
            
    except (IOError, OSError, UnicodeDecodeError):
        return False


class DICOMDirectory:
    """
    DICOMDirectory traverses a given directory to find DICOM files, and
    groups them into Series based on SeriesInstanceUID per folder.
    """

    def __init__(self, directory: str, max_workers: int = None):
        """
        :param directory: Path to the directory where DICOM files are stored.
        :param max_workers: 并行处理文件的最大线程数，默认为 CPU 核心数 * 4
        """
        self.directory = directory
        if max_workers is None:
            max_workers = min(32, (os.cpu_count() or 4) * 4)
        self.max_workers = max_workers

    def get_dicom_series(self) -> Generator["Series", None, None]:
        """
        Traverse each folder in the directory, gather files into series, and yield 
        the series for that folder. This ensures we do not load the entire directory's 
        DICOM files into memory at once.
        """
        for root, _, files in os.walk(self.directory):
            # For each folder, gather Series objects from all DICOM files in it
            series_dict = self._gather_series_from_files(root, files)

            # Yield the series for this folder before moving on
            yield from series_dict.values()

    def _process_single_file(
        self, filepath: str
    ) -> Optional[tuple[str, DICOMInstance]]:
        """
        处理单个文件，返回 (series_uid, instance) 或 None
        
        :param filepath: 文件完整路径
        :return: (series_uid, instance) 元组或 None
        """
        try:
            # 快速检查：先读取文件头判断是否为DICOM文件
            if not is_standard_dicom_file(filepath):
                return None
            
            # 确认是DICOM文件后，才创建实例并读取完整文件
            instance = DICOMInstance(filepath)
            series_uid = instance.get_tag('SeriesInstanceUID')
            if series_uid:
                return (series_uid, instance)
        except Exception:
            # 静默跳过非DICOM文件或读取失败的文件
            pass
        return None

    def _gather_series_from_files(
        self, root: str, files: List[str]
    ) -> Dict[str, "Series"]:
        """
        给定文件夹路径和文件列表，并行处理文件：
          1. 快速检查文件头判断是否为DICOM文件
          2. 只对确认的DICOM文件进行完整读取
          3. 按SeriesInstanceUID组织到字典中

        :param root: 当前扫描的文件夹路径
        :param files: 当前文件夹中的文件名列表
        :return: 该文件夹中所有文件的字典 {series_uid: Series}
        """
        series_dict: Dict[str, Series] = {}
        
        for filename in files:
            filepath = os.path.join(root, filename)
            try:
                instance = DICOMInstance(filepath)
                series_uid = instance.get_tag('SeriesInstanceUID')
                if series_uid:
                    if series_uid not in series_dict:
                        series_dict[series_uid] = Series(series_uid)
                    series_dict[series_uid].add_instance(instance)
            except Exception as exc:
                # 静默跳过错误文件
                pass

        return series_dict