"""
DICOM uploader module.

This module provides functionality for uploading DICOM files to the Orthanc server.
"""

import os
import time
import sys
from typing import Dict, Any, List, Optional, TypedDict
from concurrent.futures import ThreadPoolExecutor, as_completed, Future

import requests
from requests import RequestException, Timeout
from requests.adapters import HTTPAdapter
from pyorthanc import Orthanc


class UploadData(TypedDict):
    filepath: str


class UploadResponse(TypedDict):
    success: bool
    response: Dict[str, Any]


class DICOMUploader:
    """Class for uploading DICOM files to Orthanc server."""

    def __init__(self, config):
        """
        Initialize the DICOM uploader.

        Args:
            config: Configuration namespace with orthanc settings
        """

        self.config = config
        self.session = self._init_session()
        # Normalize base_url and compute the correct instances endpoint.
        # The final endpoint should look like: http://host:port/api/v1/orthanc/instances/
        base_url = str(config.base_url).rstrip('/')
        # Ensure the base path for API exists
        if '/api/v1/orthanc' not in base_url:
            # Assuming a base like http://localhost:29999, append the standard API path
            base_url = f"{base_url}/api/v1/orthanc"

        # Construct the instances endpoint, ensuring it ends with /instances/
        self.instances_endpoint = f"{base_url.split('/instances')[0].rstrip('/')}/instances/"

        self.DEFAULT_CONNECT_TIMEOUT = int(config.DEFAULT_CONNECT_TIMEOUT)
        self.DEFAULT_READ_TIMEOUT = int(config.DEFAULT_READ_TIMEOUT)
        self.DEFAULT_RETRY_DELAY = int(config.DEFAULT_RETRY_DELAY)
        self.DEFAULT_BATCH_SIZE = int(config.DEFAULT_BATCH_SIZE)

    def _init_session(self) -> requests.Session:
        """
        Create and configure the requests session.

        Returns:
            requests.Session: Configured session.
        """
        session = requests.Session()
        # Put cookie into session headers (if provided). Keep key lowercase as in config.
        cookie_value = getattr(self.config, 'cookie', None)
        if cookie_value:
            # Keep Cookie header for compatibility
            session.headers.update({"Cookie": cookie_value})
            # Also populate session.cookie_jar for requests to send cookies properly.
            # cookie_value may be like 'ls=xxxxx' or 'k1=v1; k2=v2'
            try:
                parts = [p.strip() for p in cookie_value.split(';') if p.strip()]
                for part in parts:
                    if '=' in part:
                        k, v = part.split('=', 1)
                        session.cookies.set(k.strip(), v.strip())
            except Exception:
                # If parsing fails, keep header-only approach.
                pass
        # 优化连接池大小：pool_maxsize 应该大于 max_workers，以支持更多并发连接
        # 这样可以充分利用 HTTP 连接复用，提高上传速度
        pool_size = max(self.config.max_workers * 2, 50)
        adapter = HTTPAdapter(
            max_retries=self.config.max_retries,
            pool_connections=self.config.max_workers,
            pool_maxsize=pool_size,  # 增加连接池大小以支持更多并发连接
        )
        session.mount("http://", adapter)
        session.mount("https://", adapter)
        return session

    def validate(self, data: UploadData) -> bool:
        """
        Validate the data before upload.

        Args:
            data: Dictionary containing the file path.

        Returns:
            bool: True if validation passes, False otherwise.
        """
        filepath = data.get("filepath")
        if not filepath or not isinstance(filepath, str):
            return False
        return os.path.exists(filepath)

    def upload(self, data: UploadData) -> UploadResponse:
        """
        Perform the actual upload request.

        Args:
            data: Dictionary containing the file path.

        Returns:
            UploadResponse: Dictionary containing success flag and response data.

        Raises:
            ValueError: If validation fails.
            Exception: If upload fails with a non-200 status.
        """
        if not self.validate(data):
            raise ValueError(f"Invalid data: {data}")

        filepath = data["filepath"].replace(r'\\', r'/')

        try:
            # 优化文件上传：使用流式上传，避免一次性加载整个文件到内存
            # 对于大文件，这可以显著减少内存使用并提高上传速度
            with open(filepath, "rb") as f:
                # 使用文件名而不是完整路径，减少传输数据量
                filename = os.path.basename(filepath)
                files = {'file': (filename, f, 'application/octet-stream')}
                # Let requests set the Content-Type header for multipart/form-data automatically.
                # Only set the Accept header.
                headers = {"Accept": "application/json, text/plain, */*"}
                endpoint = self.instances_endpoint
                # 移除调试输出以提高性能（仅在需要时启用）
                # print(f"DEBUG: Uploading {filepath} to {endpoint}", file=sys.stderr)
                response = self.session.post(
                    endpoint,
                    files=files,
                    headers=headers,
                    timeout=(self.DEFAULT_CONNECT_TIMEOUT, self.DEFAULT_READ_TIMEOUT),
                    stream=False  # 对于文件上传，stream=False 通常更快
                )

            # Orthanc may return 200 or 201 depending on the API/version; accept both as success.
            if response.status_code in (200, 201):
                # Check if response content is empty
                if not response.text:
                    return {"success": True, "response": {}}

                # Try to parse JSON
                try:
                    json_response = response.json()
                    return {"success": True, "response": json_response}
                except ValueError:
                    # If cannot parse JSON, return raw text
                    return {"success": True, "response": {"text": response.text}}

            # Non-200 status code is considered a failed upload
            # 仅在失败时输出调试信息
            error_msg = f"Upload failed with status {response.status_code}: {response.text[:200]}"
            print(f"ERROR: {error_msg}", file=sys.stderr)
            raise Exception(error_msg)

        except Exception as e:
            # Re-raise while preserving stack trace; additional logging can be added here.
            raise

    def _upload_single_file(self, filepath: str) -> bool:
        """
        Upload a single DICOM file without retry logic.

        Args:
            filepath: Path to the DICOM file.

        Returns:
            bool: True if upload successful, False otherwise.

        Raises:
            FileNotFoundError: If the file does not exist.
            Exception: If the upload fails.
        """
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"File not found: {filepath}")

        result = self.upload({"filepath": filepath})
        return result["success"]

    def _upload_file_with_retry(self, filepath: str, attempt: int = 1) -> bool:
        """
        Upload a single DICOM file with retry logic.

        Args:
            filepath: Path to the DICOM file.
            attempt: Current attempt number.

        Returns:
            bool: True if upload is successful, False otherwise.

        Raises:
            Exception: If the upload fails after all retries.
        """
        try:
            return self._upload_single_file(filepath)

        except (Timeout, RequestException) as e:
            if attempt < self.config.max_retries:
                time.sleep(self.DEFAULT_RETRY_DELAY)
                return self._upload_file_with_retry(filepath, attempt + 1)
            raise Exception(f"Upload failed after {attempt} attempts: {str(e)}")

        except Exception as e:
            if attempt < self.config.max_retries:
                time.sleep(self.DEFAULT_RETRY_DELAY)
                return self._upload_file_with_retry(filepath, attempt + 1)
            raise Exception(f"Upload failed after {attempt} attempts: {str(e)}")

    def upload_file(self, filepath: str) -> bool:
        """
        Public method to upload a single DICOM file with retries.

        Args:
            filepath: Path to the DICOM file.

        Returns:
            bool: True if upload is successful, False otherwise.

        Raises:
            FileNotFoundError: If file doesn't exist.
            Exception: If upload fails.
        """
        return self._upload_file_with_retry(filepath)

    def upload_series(self, series) -> List[bool]:
        """
        Upload all DICOM files in a series using parallel processing.

        Args:
            series: Series object containing DICOM instances.

        Returns:
            List[bool]: A list indicating success for each file.

        Raises:
            Exception: If critical upload failures occur.
        """
        from tqdm import tqdm

        total_files = len(series.instances)
        if total_files == 0:
            return []

        # 预先收集所有文件路径，避免重复访问
        filepaths = [instance.filepath for instance in series.instances]
        
        results = [False] * total_files  # 预分配结果列表
        
        # Create progress bar for file uploads
        with tqdm(total=total_files, desc="上传DICOM文件", unit="个", ncols=80) as pbar:
            # 一次性提交所有任务，让线程池管理并发，提高上传速度
            with ThreadPoolExecutor(max_workers=self.config.max_workers) as executor:
                # 一次性提交所有任务，让线程池管理并发
                futures: Dict[Future[bool], int] = {
                    executor.submit(self.upload_file, filepath): idx
                    for idx, filepath in enumerate(filepaths)
                }

                # 处理完成的任务
                success_count = 0
                failure_count = 0
                completed_count = 0
                # 批量更新进度条，减少 IO 操作频率，提高性能
                update_interval = max(1, total_files // 100)  # 根据文件数量动态调整更新间隔
                
                for future in as_completed(futures):
                    idx = futures[future]
                    try:
                        success = future.result()
                        results[idx] = success
                        completed_count += 1
                        
                        if success:
                            success_count += 1
                        else:
                            failure_count += 1
                        
                        # 优化进度条更新：减少更新频率以提高性能
                        if completed_count % update_interval == 0 or completed_count == total_files:
                            pbar.set_description(f"上传DICOM文件 [成功:{success_count}/{completed_count}]")
                        pbar.update(1)

                    except Exception as e:
                        results[idx] = False
                        completed_count += 1
                        failure_count += 1
                        
                        # 优化进度条更新频率
                        if completed_count % update_interval == 0 or completed_count == total_files:
                            pbar.set_description(f"上传DICOM文件 [失败:{failure_count}/{completed_count}]")
                        pbar.update(1)

                        # Check if we should abort (failure rate > 50%)
                        if failure_count > 0 and failure_count / completed_count > 0.5:
                            pbar.close()
                            raise Exception(
                                f"上传失败: 超过50%的文件失败 "
                                f"({failure_count}/{completed_count})"
                            )

        successful_uploads = sum(1 for r in results if r)
        print(f"\n上传完成: {successful_uploads}/{total_files} 个文件成功", file=sys.stderr)

        return results
