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

"""
export PROJECT_NAME = pull_port # 项目名
export RUN_NAME = run_1111 # 任务名
export SAVE_INTERVAL = 5000 # 模型存储最小间隔， >1000
export DATASET_DIR = /datasets/pull_port8_data/libero # 数据集路径
export BATCH_SIZE = 32 # batch_size 最大32
export NUM_TRAIN_STEPS = 20000 # 训练步长，<50000
export LOG_INTERVAL = 10 # log打印间隔，< 100
export WANDB_ENABLED = true # 是否启动wandb
export WANDB_API_KEY = "" # 如果WANDB_ENABLED为true，这里填key
export PRETRAIN_MODEL = "" # 预训练模型目录
export RESUME = true # 如果为true，需要有可恢复的现场
export OUTPUT_DIR = "" # gr00t模型存储的路径为OUTPUT_DIR, actibot_vla为 f"{OUTPUT_DIR}/{PROJECT_NAME}/{RUN_NAME}"
"""


def safe_extract_zip(zip_path: str, extract_dir: str, strip_top_level: bool = True) -> None:
    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:
    _upload_executor: Optional[concurrent.futures.ThreadPoolExecutor] = None
    _upload_executor_lock = threading.Lock()
    _upload_pending_slots: Optional[threading.BoundedSemaphore] = None

    def __init__(self) -> None:
        connect_timeout = int(os.getenv("S3_CONNECT_TIMEOUT", "10"))
        read_timeout = int(os.getenv("S3_READ_TIMEOUT", "600"))
        max_attempts = int(os.getenv("S3_MAX_ATTEMPTS", "5"))
        self.s3_client = boto3.client(
            "s3",
            endpoint_url=os.getenv("S3_ENDPOINT_URL"),  # 必须写完整的 http/https 协议和端口
            aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID", "robodock"),
            aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY", "robodock"),
            config=Config(
                connect_timeout=connect_timeout,
                read_timeout=read_timeout,
                retries={"max_attempts": max_attempts, "mode": "standard"},
            ),  # MinIO 推荐使用 v4 签名
            region_name=os.getenv("S3_REGION_NAME", "auto"),  
        )
        self.bucket_name = "robodock"
        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 = os.getenv("ROBOCK_SERVER_URL", "http://robodock-service:80")
        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:
            # ThreadPoolExecutor.shutdown 在不同 Python 版本中签名可能不同；这里保持兼容。
            executor.shutdown(wait=True)
    
    def get_dataset_info(self, uuid: str) -> dict:
        """
        get_dataset_info 的 Docstring
        
        :param self: 说明
        :param uuid: 说明
        :type uuid: str
        :return: 说明
        :rtype: dict[Any, Any]
        """
        response = requests.get(f"{self.server_url}/{self.api_dataset_prefix}?uuid={uuid}")
        if response.status_code != 200:
            raise Exception(f"Failed to get data info for {uuid}: {response.text}") 
        return response.json()

    def get_singledata_info(self, uuid: str) -> dict:
        """
        get_singledata_info 的 Docstring
        
        :param self: 说明
        :param uuid: 说明
        :type uuid: str
        :return: 说明
        :rtype: dict[Any, Any]
        """
        response = requests.get(f"{self.server_url}/{self.api_singledata_prefix}/{uuid}")
        if response.status_code != 200:
            raise Exception(f"Failed to get data info for {uuid}: {response.text}") 
        return response.json()
    
    def get_run_info(self, unique_name: str, unique_type: str = "uuid") -> dict:
        """
        get_singledata_info 的 Docstring
        
        :param self: 说明
        :param uuid: 说明
        :type uuid: str
        :return: 说明
        :rtype: dict[Any, Any]
        """
        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 data info for {unique_name}: {response.text}") 
        return response.json()
        
    def download_data(self, data_name: str, uuid: str, dirname: str) -> str:
        """
        download_data 的 Docstring
        
        :param self: 说明
        :param data_name: 说明
        :type data_name: str
        :param uuid: 说明
        :type uuid: str
        :return: 解压后的目录路径
        :rtype: str
        """
        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_check_point(self, online_path: str, local_path: str) -> str:
        """
        download_check_point 的 Docstring
        
        :param self: 说明
        :param online_path: S3上的文件路径
        :type online_path: str
        :param local_path: 本地解压目录
        :type local_path: str
        :return: 解压后的路径
        :rtype: str
        """
        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:
        """Compress the output directory into a tar and upload it to S3.

        The S3 object key will be `upload_path/<outdir_name>.tar`.
        """
        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}"

        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:
            # cleanup temporary 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]]:
        """Upload data asynchronously with bounded concurrency/queue.

        不会阻塞调用者；当排队已满时直接跳过（返回 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:
            # should not happen, but keep behavior safe
            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]]:
        """Upload a checkpoint directory in background and update run on success."""

        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:
        """
        update_singledata 的 Docstring
        
        :param self: 说明
        :param uuid: 说明
        :type uuid: str
        :param checker_result: 说明
        :type checker_result: str
        :param checker_info: 说明
        :type checker_info: str
        """
        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:
        """
        update_run 的 Docstring
        
        :param self: 说明
        :param unique_name: 说明
        :type unique_name: str
        :param body: 说明
        :type body: dict
        :return: 说明
        :rtype: int
        """
        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

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    dataset_uuid = os.getenv("DATASET_UUID", "")
    logging.info(f"uuids: {dataset_uuid}")
    robodock_action = RoboDockActions()
    env = os.environ.copy()
    logging.info(f"Environment variables: {env}")
    project_name = env.get("PROJECT_NAME", "default_project")
    is_resume = env.get("RESUME", "false").lower() == "true"
    run_name = env.get("RUN_NAME", "default_run")
    resume_from = env.get("RESUME_FROM", "")
    unique_name = f"{project_name}-{run_name}"
    outdir_path=f"{env.get('OUTPUT_DIR', '/workspace/checkpoints')}/{project_name}/{run_name}"
    try:
        # Download data
        robodock_action.update_run(unique_name, 
                                   {"status": "processing",
                                    "start_time": int(time.time() * 1000),
                                    "finish_time": None}, 
                                   unique_type="unique_name")
        data_info = robodock_action.get_dataset_info(dataset_uuid).get("data", {})
        if len(data_info) != 1:
            raise Exception(f"Dataset info invalid: {data_info}")
        data_uuids = data_info[0]['linked_singledata_list'] # 这里目前只考虑parquet数据，数据集与数据一一关联
        if not data_uuids:
            raise Exception("No linked singledata found in dataset.")
        if len(data_uuids) > 1:
            logging.warning("Multiple linked singledata found, using the first one.")
        singledata_uuid = data_uuids[0]
        singledata_info = robodock_action.get_singledata_info(singledata_uuid).get("data", {})
        data_name = singledata_info.get("data_path", "").split("/")[-1] # TODO: 兼容更多数据格式
        local_data_path = robodock_action.download_data(data_name, singledata_uuid, dirname="singledata")
        logging.info(f"Local dataset path set to: {local_data_path}")
    except Exception as e:
        error_info = f"Error during robodock action download: {e}"
        logging.error(error_info)
        robodock_action.update_run(unique_name, 
                                   {"status": "failed", 
                                    "fail_info": error_info, 
                                    "finish_time": int(time.time() * 1000)}, 
                                    unique_type="unique_name")
        raise e
    if is_resume:
        # Download latest checkpoint from robodock storage
        try:
            run_info = robodock_action.get_run_info(resume_from).get("data", {})
            check_points = run_info.get("check_points_path", [])
            if not check_points or not isinstance(check_points, list):
                raise Exception("No checkpoint found for resume.")
            last_check_point = check_points[-1]
            logging.info("Resuming training from previous checkpoint.")
            robodock_action.download_check_point(
                online_path=last_check_point,
                local_path=outdir_path
            )
            logging.info("Checkpoint download completed successfully.")
        except Exception as e:
            error_info = f"Error during checkpoint download for resume: {e}"
            logging.error(error_info)
            robodock_action.update_run(unique_name, 
                                       {"status": "failed", 
                                        "fail_info": error_info, 
                                        "finish_time": int(time.time() * 1000)}, 
                                       unique_type="unique_name")
            raise e
    try:
        # Here would be the training logic
        logging.info("Starting training process...")
        env["DATASET_DIR"] = local_data_path
        run_type = "resume" if is_resume else "no-resume"
        cmd = (
            "python scripts/gr00t_finetune.py "
            f"--dataset-path {local_data_path} "
            f"--max-steps {env.get('NUM_TRAIN_STEPS', '10000')} "
            f"--save-steps {env.get('SAVE_INTERVAL', '1000')} "
            f"--logging_steps {env.get('LOG_INTERVAL', '100')} "
            f"--output_dir {outdir_path} "
            "--num-gpus 1 "
            "--data-config custom_data_config:FiiWR1T1DataConfig "
            "--report_to azure_ml "
            "--dataloader_num_workers 2 "
            f"--{run_type}"
        )
        
        subprocess.run(["bash", "-lc", cmd], check=True, env=env)
        logging.info("Training completed successfully.")
        # robodock_action.wait_for_uploads()
        # Update robodock with success status
    except Exception as e:
        error_info = f"Error during training process: {e}"
        logging.error(error_info)
        # Update robodock with failed status
        robodock_action.update_run(unique_name, 
                                   {"status": "failed", 
                                    "fail_info": error_info, 
                                    "finish_time": int(time.time() * 1000)}, 
                                    unique_type="unique_name")
        raise e
    try:
        # upload output model files will be in train process
        logging.info("Starting model upload process...")
        # robodock_action.upload_data(
        #     outdir_path=f"{env.get('OUTPUT_DIR', '/workspace/checkpoints')}/{project_name}/{run_name}",
        #     upload_path=f"models/{project_name}/{run_name}")
        robodock_action.update_run(unique_name, 
                                   {"status": "success", 
                                    "finish_time": int(time.time() * 1000)}, 
                                    unique_type="unique_name")
        logging.info("Model upload completed successfully.")
        
        # 删除本地 checkpoint 文件夹以释放空间
        outdir_path = f"{env.get('OUTPUT_DIR', '/workspace/checkpoints')}/{project_name}/{run_name}"
        if Path(outdir_path).exists():
            shutil.rmtree(outdir_path)
            logging.info(f"Removed local checkpoint directory: {outdir_path}")
    
    except Exception as e:
        error_info = f"Error during model upload process: {e}"
        logging.error(error_info)
        robodock_action.update_run(unique_name, 
                                   {"status": "failed", 
                                    "fail_info": error_info, 
                                    "finish_time": int(time.time() * 1000)}, 
                                    unique_type="unique_name")
        raise e
    
    