import docker
import os
import logging
from collections import deque


class DockerRunner:
    def __init__(self, project_path: str, config: dict):
        """
        初始化 DockerRunner

        Args:
            project_path (str): 项目路径，将挂载到 Docker 容器中
        """
        self.project_path = project_path
        self.config = config
        self.logger = logging.getLogger("UAV_Preprocess.DockerRunner")
        self.docker_client = docker.from_env()

    def run_odm_container(self):
        """
        使用 Docker SDK 运行 OpenDroneMap 容器
        """
        try:
            self.logger.info("开始运行docker run指令")
            # 挂载路径
            volume_mapping = {
                self.project_path: {
                    'bind': '/datasets',
                    'mode': 'rw'
                }
            }

            # Docker 命令参数
            command = [
                "--project-path", "/datasets",
                "project",
                "--max-concurrency", "15",
                "--force-gps",
                "--split-overlap", "0",
                "--rerun-all"
            ]

            if self.config.mode == "快拼模式":
                command += ["--fast-orthophoto", "--skip-3dmodel"]
            else:
                command += ["--dsm", "--dtm"]

            # 运行容器
            container = self.docker_client.containers.run(
                image="opendronemap/odm:gpu",
                command=command,
                volumes=volume_mapping,
                device_requests=[
                    docker.types.DeviceRequest(
                        count=-1, capabilities=[["gpu"]])
                ],  # 添加 GPU 支持
                remove=False,  # 容器运行结束后不自动删除，便于获取日志
                tty=True,
                detach=True  # 后台运行
            )

            # 等待容器运行完成
            exit_status = container.wait()
            if exit_status["StatusCode"] != 0:
                self.logger.error(f"容器运行失败，退出状态码: {exit_status['StatusCode']}")

                # 获取容器的错误日志
                error_logs = container.logs(
                    stderr=True).decode("utf-8").splitlines()
                self.logger.error("容器运行失败的详细错误日志：")
                for line in error_logs:
                    self.logger.error(line)

            else:
                # 获取所有日志
                logs = container.logs().decode("utf-8").splitlines()

                # 输出最后 50 行日志
                self.logger.info("容器运行完成，以下是最后 50 行日志：")
                for line in logs[-50:]:
                    self.logger.info(line)

            # 删除容器
            container.remove()

        except Exception as e:
            self.logger.error(f"运行 Docker 容器时发生错误: {str(e)}", exc_info=True)
            raise


if __name__ == "__main__":
    # 示例用法
    project_path = r"E:\datasets\UAV\199"
    docker_runner = DockerRunner(project_path)
    docker_runner.run_odm_container()