"""RoboDock Actions - Core functionality for data and model management."""

import boto3
import requests
import logging
import threading
import concurrent.futures
from botocore.config import Config
from typing import Any, Optional, Callable
import zipfile
import tarfile
import os
import time
from pathlib import Path


def safe_extract_zip(zip_path: str, extract_dir: str, strip_top_level: bool = True) -> None:
    """安全解压ZIP文件，防止zip-slip攻击。
    
    Args:
        zip_path: ZIP文件路径
        extract_dir: 解压目标目录
        strip_top_level: 是否移除顶层目录
    """
    extract_dir_new = Path(extract_dir).resolve()
    extract_dir_new.mkdir(parents=True, exist_ok=True)

    with zipfile.ZipFile(zip_path, 'r') as z:
        # 列出所有文件名（排除目录条目）
        names = [n for n in z.namelist() if n and not n.endswith('/')]
        # 检测是否存在统一的顶层目录名
        prefix = None
        if strip_top_level and names:
            first_parts = [Path(n).parts[0] for n in names if Path(n).parts]
            if first_parts and all(p == first_parts[0] for p in first_parts):
                prefix = first_parts[0] + '/'

        for member in z.infolist():
            name = member.filename
            if prefix and name.startswith(prefix):
                rel = name[len(prefix):]
            else:
                rel = name
            if not rel:
                continue  # 跳过根目录条目
            target_path = (extract_dir_new / rel).resolve()

            # 安全检查：防 zip-slip（确保目标路径在 extract_dir 下）
            if not str(target_path).startswith(str(extract_dir_new) + os.sep):
                raise Exception(f"Unsafe zip member path: {name}")

            target_path.parent.mkdir(parents=True, exist_ok=True)
            if member.is_dir():
                continue
            with z.open(member) as src, open(target_path, "wb") as dst:
                dst.write(src.read())


class RoboDockActions:
    """RoboDock核心操作类，提供数据上传下载、任务管理等功能。"""

    _upload_executor: Optional[concurrent.futures.ThreadPoolExecutor] = None
    _upload_executor_lock = threading.Lock()
    _upload_pending_slots: Optional[threading.BoundedSemaphore] = None

    def __init__(
        self,
        s3_endpoint: Optional[str] = None,
        s3_access_key: Optional[str] = None,
        s3_secret_key: Optional[str] = None,
        api_url: Optional[str] = None,
        bucket_name: Optional[str] = None,
        region_name: Optional[str] = None,
    ) -> None:
        """初始化RoboDock Actions客户端。
        
        Args:
            s3_endpoint: S3存储服务端点
            s3_access_key: S3访问密钥
            s3_secret_key: S3密钥
            api_url: RoboDock API URL
            bucket_name: S3存储桶名称
            region_name: S3区域名称
        """
        # 优先使用入参；未传入时回退到环境变量；环境变量未设置则用兜底值
        s3_endpoint = s3_endpoint or os.getenv("S3_ENDPOINT_URL", "http://192.168.1.6:9000")
        s3_access_key = s3_access_key or os.getenv("AWS_ACCESS_KEY_ID", "robodock")
        s3_secret_key = s3_secret_key or os.getenv("AWS_SECRET_ACCESS_KEY", "robodock")
        api_url = api_url or os.getenv("ROBOCK_SERVER_URL", "http://robodock-service:80")
        bucket_name = bucket_name or os.getenv("S3_BUCKET_NAME", "robodock")
        region_name = region_name or os.getenv("S3_REGION_NAME", "auto")

        self.s3_client = boto3.client(
            "s3",
            endpoint_url=s3_endpoint,
            aws_access_key_id=s3_access_key,
            aws_secret_access_key=s3_secret_key,
            config=Config(
                connect_timeout=10,
                read_timeout=600,
                retries={"max_attempts": 3, "mode": "standard"},
                signature_version="s3v4",
            ),  # MinIO 推荐使用 v4 签名
            region_name=region_name,
        )
        self.bucket_name = bucket_name
        self.api_singledata_prefix = "api/p/v1/singledata"
        self.api_run_prefix = "api/p/v1/run"
        self.api_dataset_prefix = "api/p/v1/dataset"
        self.server_url = api_url
        
        # 确保bucket存在
        try:
            self.s3_client.head_bucket(Bucket=self.bucket_name)
        except:
            self.s3_client.create_bucket(Bucket=self.bucket_name)

    def _default_upload_success(self, s3_key: str, *, outdir_path: str, upload_path: str) -> None:
        logging.info(
            "Background upload success: local=%s upload_path=%s s3_key=%s",
            outdir_path,
            upload_path,
            s3_key,
        )

    def _default_upload_error(self, exc: Exception, *, outdir_path: str, upload_path: str) -> None:
        logging.exception(
            "Background upload failed: local=%s upload_path=%s err=%s",
            outdir_path,
            upload_path,
            exc,
        )

    @classmethod
    def _get_upload_executor(cls) -> concurrent.futures.ThreadPoolExecutor:
        """Global upload executor with bounded pending queue.

        - `UPLOAD_MAX_WORKERS`: 并发上传线程数
        - `UPLOAD_MAX_PENDING`: 允许排队的任务数（包含正在执行的）
        """
        with cls._upload_executor_lock:
            if cls._upload_executor is None:
                max_workers = int(os.getenv("UPLOAD_MAX_WORKERS", "4"))
                max_pending = int(os.getenv("UPLOAD_MAX_PENDING", "8"))
                max_workers = max(1, max_workers)
                max_pending = max(max_workers, max_pending)

                cls._upload_executor = concurrent.futures.ThreadPoolExecutor(
                    max_workers=max_workers,
                    thread_name_prefix="robodock-upload",
                )
                cls._upload_pending_slots = threading.BoundedSemaphore(value=max_pending)
            return cls._upload_executor

    @classmethod
    def wait_for_uploads(cls, timeout: Optional[float] = None) -> None:
        """Wait for all queued upload tasks.

        Note: `timeout` applies to executor shutdown; if tasks hang beyond S3 timeouts,
        this can still block.
        """
        with cls._upload_executor_lock:
            executor = cls._upload_executor
            cls._upload_executor = None
            cls._upload_pending_slots = None
        if executor is not None:
            executor.shutdown(wait=True)

    def get_dataset_info(self, uuid: str) -> dict:
        """获取数据集信息。
        
        Args:
            uuid: 数据集UUID
            
        Returns:
            数据集信息字典
            
        Raises:
            Exception: 获取数据集信息失败
        """
        response = requests.get(f"{self.server_url}/{self.api_dataset_prefix}?uuid={uuid}")
        if response.status_code != 200:
            raise Exception(f"Failed to get dataset info for {uuid}: {response.text}")
        return response.json()

    def get_singledata_info(self, uuid: str) -> dict:
        """获取单条数据信息。
        
        Args:
            uuid: 数据UUID
            
        Returns:
            数据信息字典
            
        Raises:
            Exception: 获取数据信息失败
        """
        response = requests.get(f"{self.server_url}/{self.api_singledata_prefix}/{uuid}")
        if response.status_code != 200:
            raise Exception(f"Failed to get singledata info for {uuid}: {response.text}")
        return response.json()

    def get_run_info(self, unique_name: str, unique_type: str = "uuid") -> dict:
        """获取运行任务信息。
        
        Args:
            unique_name: 任务唯一标识
            unique_type: 标识类型，可选"uuid"或"unique_name"
            
        Returns:
            任务信息字典
            
        Raises:
            Exception: 获取任务信息失败
        """
        response = requests.get(
            f"{self.server_url}/{self.api_run_prefix}/{unique_name}?unique_type={unique_type}"
        )
        if response.status_code != 200:
            raise Exception(f"Failed to get run info for {unique_name}: {response.text}")
        return response.json()

    def download_data(self, data_name: str, uuid: str, dirname: str) -> str:
        """从S3下载并解压数据。
        
        Args:
            data_name: 数据文件名
            uuid: 数据UUID
            dirname: S3上的目录名
            
        Returns:
            解压后的本地目录路径
            
        Raises:
            Exception: 不支持的文件类型或下载失败
        """
        tmp_dir = Path.cwd() / "tmp"
        tmp_dir.mkdir(parents=True, exist_ok=True)
        download_path = str(tmp_dir / f"{uuid}_{data_name}")
        extract_path = str(tmp_dir / uuid)

        self.s3_client.download_file(self.bucket_name, f"{dirname}/{data_name}", download_path)
        logging.info(f"Downloaded {data_name} to {download_path}")

        # 根据文件扩展名解压
        if data_name.endswith('.zip'):
            safe_extract_zip(download_path, extract_path)
            logging.info(f"Extracted zip to {extract_path}")
        elif data_name.endswith('.tar.gz') or data_name.endswith('.tgz'):
            with tarfile.open(download_path, 'r:gz') as tar:
                tar.extractall(path=extract_path)
            logging.info(f"Extracted tar.gz to {extract_path}")
            os.remove(download_path)
        elif data_name.endswith('.tar'):
            with tarfile.open(download_path, 'r') as tar:
                tar.extractall(path=extract_path)
            logging.info(f"Extracted tar to {extract_path}")
            os.remove(download_path)
        else:
            raise Exception(
                f"Unsupported file type: {data_name}. Supported: .zip, .tar, .tar.gz, .tgz"
            )

        return extract_path

    def download_checkpoint(self, online_path: str, local_path: str) -> str:
        """从S3下载检查点文件并解压。
        
        Args:
            online_path: S3上的文件路径
            local_path: 本地解压目录
            
        Returns:
            解压后的路径
            
        Raises:
            Exception: 不支持的文件类型
        """
        filename = online_path.split('/')[-1]
        download_path = f"{local_path}/{filename}"

        # 确保本地目录存在
        Path(local_path).mkdir(parents=True, exist_ok=True)

        self.s3_client.download_file(self.bucket_name, online_path, download_path)
        logging.info(f"Downloaded {filename} to {download_path}")

        # 解压文件
        if filename.endswith('.zip'):
            extract_path = download_path.rstrip('.zip')
            safe_extract_zip(download_path, extract_path)
            logging.info(f"Extracted zip to {extract_path}")
            return extract_path
        elif filename.endswith('.tar.gz') or filename.endswith('.tgz'):
            with tarfile.open(download_path, 'r:gz') as tar:
                tar.extractall(path=local_path)
            logging.info(f"Extracted tar.gz to {local_path}")
            os.remove(download_path)
            return local_path
        elif filename.endswith('.tar'):
            with tarfile.open(download_path, 'r') as tar:
                tar.extractall(path=local_path)
            logging.info(f"Extracted tar to {local_path}")
            os.remove(download_path)
            return local_path
        else:
            raise Exception(f"Unsupported file type: {filename}")

    def upload_data(self, outdir_path: str, upload_path: str) -> str:
        """压缩本地目录并上传到S3。
        
        Args:
            outdir_path: 本地输出目录路径
            upload_path: S3上传路径
            
        Returns:
            S3对象键
            
        Raises:
            FileNotFoundError: 输出目录不存在
            Exception: 压缩或上传失败
        """
        outdir = Path(outdir_path)
        if not outdir.exists():
            raise FileNotFoundError(f"Output directory does not exist: {outdir}")

        tar_name = f"{outdir.name}.tar"
        tmp_tar = outdir.parent / tar_name
        s3_key = f"{upload_path.rstrip('/')}/{tar_name}".lstrip('/')

        try:
            logging.info("Creating tar %s from %s", tmp_tar, outdir)
            with tarfile.open(tmp_tar, "w") as tar:
                tar.add(outdir, arcname=outdir.name)

            logging.info("Uploading %s to s3://%s/%s", tmp_tar, self.bucket_name, s3_key)
            self.s3_client.upload_file(str(tmp_tar), self.bucket_name, s3_key)
            logging.info("Upload successful: s3://%s/%s", self.bucket_name, s3_key)
        except Exception as e:
            logging.error("Failed to create or upload tar: %s", e)
            raise
        finally:
            # 清理临时tar文件
            try:
                if tmp_tar.exists():
                    tmp_tar.unlink()
                    logging.debug("Removed temporary tar %s", tmp_tar)
            except Exception as cleanup_err:
                logging.warning("Failed to remove temporary tar %s: %s", tmp_tar, cleanup_err)
        return s3_key

    def upload_data_async(
        self,
        outdir_path: str,
        upload_path: str,
        on_success: Optional[Callable[[str], None]] = None,
        on_error: Optional[Callable[[Exception], None]] = None,
    ) -> Optional[concurrent.futures.Future[None]]:
        """异步上传数据，队列满时返回None。

        Args:
            outdir_path: 本地输出目录路径
            upload_path: S3上传路径
            on_success: 成功回调（参数为s3_key）
            on_error: 失败回调（参数为异常）

        Returns:
            Future对象或None（队列已满）
        """
        if on_success is None:
            on_success = lambda s3_key: self._default_upload_success(
                s3_key, outdir_path=outdir_path, upload_path=upload_path
            )

        if on_error is None:
            on_error = lambda exc: self._default_upload_error(
                exc, outdir_path=outdir_path, upload_path=upload_path
            )

        executor = self._get_upload_executor()
        slots = self.__class__._upload_pending_slots
        if slots is None:
            slots = threading.BoundedSemaphore(value=1)
            self.__class__._upload_pending_slots = slots

        if not slots.acquire(blocking=False):
            logging.warning(
                "Upload queue full; skip upload: local=%s upload_path=%s",
                outdir_path,
                upload_path,
            )
            return None

        def _run() -> None:
            try:
                s3_key = self.upload_data(outdir_path=outdir_path, upload_path=upload_path)
                on_success(s3_key)
            except Exception as exc:
                on_error(exc)
            finally:
                try:
                    slots.release()
                except Exception:
                    pass

        return executor.submit(_run)

    def upload_checkpoint_async(
        self,
        *,
        outdir_path: str,
        upload_path: str,
        run_unique_name: str,
        unique_type: str = "unique_name",
        run_field: str = "check_point_path",
    ) -> Optional[concurrent.futures.Future[None]]:
        """异步上传checkpoint目录，并在成功后更新run。"""

        def _on_success(s3_key: str) -> None:
            self._default_upload_success(s3_key, outdir_path=outdir_path, upload_path=upload_path)
            self.update_run(run_unique_name, {run_field: s3_key}, unique_type=unique_type)

        def _on_error(exc: Exception) -> None:
            self._default_upload_error(exc, outdir_path=outdir_path, upload_path=upload_path)
            self.update_run(
                run_unique_name,
                {
                    "status": "failed",
                    "finish_time": int(time.time() * 1000),
                    "fail_info": str(exc),
                },
                unique_type=unique_type,
            )
            os._exit(1)

        return self.upload_data_async(
            outdir_path=outdir_path,
            upload_path=upload_path,
            on_success=_on_success,
            on_error=_on_error,
        )

    def update_singledata(self, uuid: str, checker_result: str, checker_info: str) -> int:
        """更新单条数据的检查结果。
        
        Args:
            uuid: 数据UUID
            checker_result: 检查结果
            checker_info: 检查信息
            
        Returns:
            HTTP状态码
            
        Raises:
            Exception: 更新失败
        """
        data = {
            "any_in_dataformat": {
                "checker_info": checker_info,
                "checker_result": checker_result
            }
        }
        response = requests.patch(
            f"{self.server_url}/{self.api_singledata_prefix}/{uuid}",
            json=data
        )
        logging.info(f"Update result: {response.status_code}, {response.text}")
        if response.status_code != 200:
            logging.error(f"Failed to update singledata for {uuid}: {response.text}")
            raise Exception(f"Failed to update singledata for {uuid}: {response.text}")
        return response.status_code

    def update_run(self, unique_name: str, body: dict, unique_type: str = "uuid") -> int:
        """更新运行任务状态。
        
        Args:
            unique_name: 任务唯一标识
            body: 更新内容字典
            unique_type: 标识类型
            
        Returns:
            HTTP状态码
            
        Raises:
            Exception: 更新失败
        """
        logging.info(f"Updating run {unique_name} with body: {body}")
        response = requests.patch(
            f"{self.server_url}/{self.api_run_prefix}/{unique_name}?unique_type={unique_type}",
            json=body
        )
        logging.info(f"Update result: {response.status_code}, {response.text}")
        if response.status_code != 200:
            logging.error(f"Failed to update run for {unique_name}: {response.text}")
            raise Exception(f"Failed to update run for {unique_name}: {response.text}")
        return response.status_code
