"""Docker Engine API over a Unix socket or a pinned SSH connection.

Only the constant `docker system dial-stdio` command runs on the server.
All variable data travels as JSON, never interpolated into a shell command.
"""

from __future__ import annotations

import base64
import hashlib
import http.client
import json
import posixpath
import socket
from urllib.parse import quote

from .settings import PoolSettings


class DockerError(RuntimeError):
    pass


_WORKER_STORAGE_SELECTION = """\
import os
from pathlib import Path
import yaml

config_path = Path(os.getenv('DEER_FLOW_CONFIG_PATH') or 'config.yaml')
config = yaml.safe_load(config_path.read_text(encoding='utf-8')) or {}
checkpoint = config.get('checkpointer')
if checkpoint is not None:
    checkpoint_type = checkpoint.get('type')
    if checkpoint_type == 'postgres':
        raw = os.getenv('DEER_FLOW_USE_PG_CHECKPOINTER')
        if raw:
            postgres_enabled = raw.strip().lower() not in {'0', 'false', 'no', 'off'}
        else:
            postgres_enabled = bool(checkpoint.get('use_postgres', True))
        backend = 'postgres' if postgres_enabled else 'sqlite'
    elif checkpoint_type == 'sqlite':
        backend = 'sqlite'
    else:
        backend = 'memory'
else:
    backend = (config.get('database') or {}).get('backend', 'memory')
    if backend not in {'postgres', 'sqlite'}:
        backend = 'memory'
"""
_WORKER_STORAGE_PROBE = _WORKER_STORAGE_SELECTION + """\
print('__DEERFLOW_STORAGE_BACKEND__=' + backend)
"""
_WORKER_PYTHON_PROBE = _WORKER_STORAGE_SELECTION + """\
import importlib.util
import os
import sys

if sys.version_info < (3, 12):
    raise SystemExit(78)
required = ['uvicorn']
if backend == 'postgres':
    required.extend((
        'langgraph.checkpoint.postgres.aio',
        'langgraph.store.postgres.aio',
        'psycopg.rows',
        'psycopg_pool',
    ))
elif backend == 'sqlite':
    required.extend((
        'langgraph.checkpoint.sqlite.aio',
        'langgraph.store.sqlite.aio',
        'aiosqlite',
    ))
missing = []
for module in required:
    try:
        if importlib.util.find_spec(module) is None:
            missing.append(module)
    except (ImportError, ModuleNotFoundError, ValueError):
        missing.append(module)
if not Path('app/gateway/worker_pool/worker.py').is_file():
    missing.append('app.gateway.worker_pool.worker')
if missing:
    packages = (
        'langgraph-checkpoint-postgres psycopg[binary] psycopg-pool'
        if backend == 'postgres'
        else 'langgraph-checkpoint-sqlite aiosqlite'
    )
    print('__DEERFLOW_STORAGE_ERROR__=Worker Python is missing configured runtime modules: '
          + ', '.join(missing) + '; install ' + packages, flush=True)
    raise SystemExit(78)
print('__DEERFLOW_WORKER_PYTHON__=' + os.path.abspath(sys.executable))
"""
_WORKER_PYTHON_PROBE_SHELL = r'''
probe() {
  candidate="$1"
  [ -n "$candidate" ] || return 1
  output=$("$candidate" -c "$2" 2>/dev/null) || {
    case "$output" in
      *__DEERFLOW_STORAGE_ERROR__=*)
        error=${output##*__DEERFLOW_STORAGE_ERROR__=}
        printf '__DEERFLOW_STORAGE_ERROR__=%s\n' "$error"
        ;;
    esac
    return 1
  }
  case "$output" in
    *__DEERFLOW_WORKER_PYTHON__=*)
      resolved=${output##*__DEERFLOW_WORKER_PYTHON__=}
      case "$resolved" in /*) printf '%s\n' "$resolved"; return 0 ;; esac
      ;;
  esac
  return 1
}
code="$1"
shift
for candidate in "$@"; do
  if probe "$candidate" "$code"; then exit 0; fi
done
# The main API may use a Python executable outside the image's default PATH.
# Probe the interpreter actually running in this source container before PATH
# fallbacks, using the same configured persistence-adapter check above.
for process_exe in /proc/[0-9]*/exe; do
  candidate=$(readlink "$process_exe" 2>/dev/null) || continue
  case "$candidate" in
    */python|*/python[0-9]|*/python[0-9].* )
      case "$candidate" in *" (deleted)"*) continue ;; esac
      if probe "$candidate" "$code"; then exit 0; fi
      ;;
  esac
done
for candidate in python3 python3.14 python3.13 python3.12 python3.11 python; do
  if probe "$candidate" "$code"; then exit 0; fi
done
exit 1
'''


class DockerEngine:
    def __init__(self, settings: PoolSettings):
        self.settings = settings

    def request(self, method: str, path: str, body: dict | None = None, *, raw=False, max_bytes=None, json_stream=False):
        client = None
        channel = None
        # Committing a large production backend layer can take longer than
        # the ordinary metadata/exec requests. Do not turn a slow but valid
        # runtime snapshot into a permanent Worker reconciliation failure.
        request_timeout = 300 if method == "POST" and path.startswith("/commit?") else 20
        try:
            if self.settings.transport == "socket":
                channel = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
                channel.settimeout(request_timeout)
                channel.connect("/var/run/docker.sock")
            else:
                import paramiko

                expected = self.settings.host_key_sha256
                if not expected.startswith("SHA256:"):
                    raise DockerError("请先填写服务器 SSH 主机密钥指纹（SHA256:…）")

                class PinnedKey(paramiko.MissingHostKeyPolicy):
                    def missing_host_key(self, ssh, hostname, key):
                        actual = "SHA256:" + base64.b64encode(hashlib.sha256(key.asbytes()).digest()).decode().rstrip("=")
                        if actual != expected:
                            raise DockerError("SSH 主机密钥指纹不匹配，连接已拒绝")

                client = paramiko.SSHClient()
                client.set_missing_host_key_policy(PinnedKey())
                client.connect(
                    self.settings.host,
                    port=self.settings.port,
                    username=self.settings.username,
                    password=self.settings.password,
                    timeout=10,
                    auth_timeout=10,
                    banner_timeout=10,
                    look_for_keys=False,
                    allow_agent=False,
                )
                channel = client.get_transport().open_session(timeout=20)
                channel.settimeout(request_timeout)
                channel.exec_command("docker system dial-stdio")
            data = json.dumps(body).encode() if body is not None else b""
            request = f"{method} {path} HTTP/1.1\r\nHost: docker\r\nContent-Type: application/json\r\nContent-Length: {len(data)}\r\nConnection: close\r\n\r\n".encode() + data
            channel.sendall(request)
            response = http.client.HTTPResponse(channel)
            response.begin()
            if json_stream and response.status < 400:
                # Exec uses a hijacked connection that may outlive the process.
                # Stop on a complete bounded JSON result instead of waiting for EOF.
                result = b""
                while len(result) < (max_bytes or 1024 * 1024):
                    chunk = response.read1(min(65536, (max_bytes or 1024 * 1024) - len(result)))
                    if not chunk:
                        break
                    result += chunk
                    try:
                        json.loads(decode_docker_output(result))
                        break
                    except ValueError:
                        pass
            else:
                result = response.read(max_bytes) if max_bytes else response.read()
            if response.status >= 400:
                # Docker errors can contain environment values / host paths.
                raise DockerError(f"Docker 操作失败（HTTP {response.status}），请检查容器、网络和账号权限")
            return result if raw else (json.loads(result) if result else {})
        except DockerError:
            raise
        except Exception as exc:
            raise DockerError(f"服务器连接失败（{type(exc).__name__}），请检查地址、账号、密码和 Docker 权限") from None
        finally:
            if channel is not None:
                channel.close()
            if client is not None:
                client.close()

    def _exec_capture(self, container_id: str, command: list[str], working_dir: str, user: str) -> tuple[int, str]:
        created = self.request("POST", f"/containers/{quote(container_id, safe='')}/exec", {
            "AttachStdout": True,
            "AttachStderr": True,
            "Cmd": command,
            "WorkingDir": working_dir,
            "User": user,
        })
        exec_id = quote(created["Id"], safe="")
        output = self.request(
            "POST", f"/exec/{exec_id}/start", {"Detach": False, "Tty": False},
            raw=True, max_bytes=1024 * 1024,
        )
        inspected = self.request("GET", f"/exec/{exec_id}/json")
        exit_code = inspected.get("ExitCode")
        return (int(exit_code) if exit_code is not None else 1), decode_docker_output(output)

    @staticmethod
    def _python_candidates(source: dict, preferred: str) -> list[str]:
        candidates = [
            preferred,
            "/opt/venv/bin/python", "/opt/venv/bin/python3",
            "/app/.venv/bin/python", "/app/.venv/bin/python3",
            "/app/backend/.venv/bin/python", "/venv/bin/python", "/venv/bin/python3",
            "/usr/local/bin/python", "/usr/local/bin/python3",
            "/usr/local/bin/python3.14", "/usr/local/bin/python3.13",
            "/usr/local/bin/python3.12", "/usr/local/bin/python3.11",
            "/usr/bin/python", "/usr/bin/python3", "/usr/bin/python3.12",
        ]
        env = source.get("Config", {}).get("Env", [])
        path_value = next((item[5:] for item in env if item.startswith("PATH=")), "")
        for directory in path_value.split(":"):
            if directory.startswith("/"):
                candidates.extend(f"{directory.rstrip('/')}/{name}" for name in ("python3", "python"))
        return list(dict.fromkeys(path for path in candidates if path and path.startswith("/")))

    def discover_worker_python(self, source: dict, preferred: str, working_dir: str) -> str:
        """Find an interpreter inside the source image that can actually start a Worker."""
        container_id = source.get("Id")
        if not container_id:
            raise DockerError("无法读取主后端容器 ID，不能检测镜像内 Python")
        config = source.get("Config", {})
        user = config.get("User", "")
        candidates = self._python_candidates(source, preferred)
        storage_error = None
        try:
            code, output = self._exec_capture(
                container_id,
                ["/bin/sh", "-c", _WORKER_PYTHON_PROBE_SHELL, "worker-python-probe", _WORKER_PYTHON_PROBE, *candidates],
                working_dir,
                user,
            )
            if code == 0:
                for line in reversed(output.splitlines()):
                    if line.startswith("/") and "\x00" not in line:
                        return line.strip()
            for line in output.splitlines():
                prefix = "__DEERFLOW_STORAGE_ERROR__="
                if line.startswith(prefix):
                    storage_error = line[len(prefix):].strip()
        except DockerError:
            # Minimal images may not include a shell; try known Python paths directly.
            pass

        for candidate in candidates:
            try:
                code, output = self._exec_capture(
                    container_id,
                    [candidate, "-c", _WORKER_PYTHON_PROBE],
                    working_dir,
                    user,
                )
            except DockerError:
                continue
            if code == 0:
                for line in reversed(output.splitlines()):
                    prefix = "__DEERFLOW_WORKER_PYTHON__="
                    if line.startswith(prefix):
                        resolved = line[len(prefix):].strip()
                        if resolved.startswith("/") and "\x00" not in resolved:
                            return resolved
            for line in output.splitlines():
                prefix = "__DEERFLOW_STORAGE_ERROR__="
                if line.startswith(prefix):
                    storage_error = line[len(prefix):].strip()
        if storage_error:
            raise DockerError(storage_error)
        raise DockerError(
            "未找到与主服务聊天存储配置兼容的 Worker Python（要求 Python 3.12+，"
            "且能导入 uvicorn、Worker 入口及当前 SQLite/PostgreSQL 适配器）；"
            "已检查主服务常见 Python 路径和运行进程解释器，请确认 Worker 复用了主服务的镜像、"
            "环境变量、挂载目录和 Python 环境"
        )

    def configured_storage_backend(self, source: dict, preferred: str, working_dir: str) -> str:
        """Read the live config so SQLite→PostgreSQL changes invalidate probes."""
        container_id = source.get("Id")
        if not container_id:
            raise DockerError("无法读取主后端容器 ID，不能检查聊天存储配置")
        user = source.get("Config", {}).get("User", "")
        for candidate in self._python_candidates(source, preferred):
            try:
                code, output = self._exec_capture(
                    container_id,
                    [candidate, "-c", _WORKER_STORAGE_PROBE],
                    working_dir,
                    user,
                )
            except DockerError:
                continue
            if code == 0:
                for line in reversed(output.splitlines()):
                    prefix = "__DEERFLOW_STORAGE_BACKEND__="
                    if line.startswith(prefix):
                        backend = line[len(prefix):].strip()
                        if backend in {"postgres", "sqlite", "memory"}:
                            return backend
        raise DockerError(
            "无法读取主服务当前聊天存储配置；请确认 Python 路径、工作目录和主服务配置文件正确"
        )

    def snapshot_worker_runtime(self, source: dict, instance: str) -> str:
        """Snapshot the running main container's writable layer for new workers.

        Docker create from Config.Image alone loses packages installed into the
        running container after it was created. A local commit carries those
        offline-installed packages; bind/volume mounts are still cloned as-is.
        """
        container_id = source.get("Id")
        if not container_id:
            raise DockerError("无法读取主后端容器 ID，不能复用其 Python 运行环境")
        # A freshly started backend often has only a tiny writable overlay
        # (logs/runtime markers). Committing it can take minutes on an offline
        # Docker host even though the immutable image already has the runtime.
        # Keep committing larger overlays that may include offline-installed
        # Python packages.
        base_image = source.get("Image")
        writable_layer_size = source.get("SizeRw")
        if (
            isinstance(base_image, str)
            and base_image
            and isinstance(writable_layer_size, int)
            and writable_layer_size <= 128 * 1024
        ):
            return base_image
        repo = f"deerflow-worker-runtime-{instance}"
        tag = container_id[:12]
        result = self.request(
            "POST",
            f"/commit?container={quote(container_id, safe='')}&repo={quote(repo, safe='')}&tag={quote(tag, safe='')}&pause=false",
        )
        image_id = result.get("Id") if isinstance(result, dict) else None
        if not isinstance(image_id, str) or not image_id.startswith("sha256:"):
            raise DockerError("无法从主容器生成 Worker 运行镜像；请确认 Docker 账号具有容器提交权限")
        return image_id

    def host_snapshot(self, container: str):
        """Only our constant read-only collector may be executed by diagnostics."""
        if self.settings.transport == "socket":
            from .host_metrics import read_snapshot
            return read_snapshot()
        command = [self.settings.python_executable, "-m", "app.gateway.worker_pool.host_metrics"]
        created = self.request("POST", f"/containers/{quote(container, safe='')}/exec", {
            "AttachStdout": True, "AttachStderr": True, "Cmd": command,
            "WorkingDir": self.settings.working_dir,
        })
        data = self.request("POST", f"/exec/{created['Id']}/start", {"Detach": False, "Tty": False}, raw=True, max_bytes=1024 * 1024, json_stream=True)
        out = decode_docker_output(data)
        try:
            return json.loads(out)
        except ValueError:
            raise DockerError("宿主机采集器未运行成功，请检查主容器代码和只读 /proc、/sys 挂载") from None


def decode_docker_output(data: bytes) -> str:
    """Docker attaches stdout/stderr in 8-byte length-prefixed frames."""
    import struct
    if len(data) < 8 or data[0] not in (0, 1, 2) or data[1:4] != b"\x00\x00\x00":
        return data.decode("utf-8", errors="replace")
    parts = []
    offset = 0
    while offset + 8 <= len(data):
        size = struct.unpack(">I", data[offset + 4:offset + 8])[0]
        parts.append(data[offset + 8:offset + 8 + size])
        offset += 8 + size
    return b"".join(parts).decode("utf-8", errors="replace")


def worker_spec(
    source: dict,
    settings: PoolSettings,
    instance: str,
    token: str,
    control_home: str,
    *,
    runtime_image: str | None = None,
    require_shared_archive: bool = False,
) -> dict:
    """Clone an explicit runtime template, without ports, privileges or Docker access."""
    if settings.network not in source["NetworkSettings"]["Networks"]:
        raise DockerError("源容器必须连接所选 Docker 网络")
    if settings.network in {"host", "none", "bridge"}:
        raise DockerError("请选择用户自定义 Docker 网络")
    mounts = []
    for mount in source.get("Mounts", []):
        target = mount["Destination"].rstrip("/")
        if control_home == target:
            continue  # dedicated controller volume is never passed to workers
        if "docker.sock" in target or target in {"/", "/var/run", "/run"}:
            continue
        if target in {"/host/proc", "/host/sys"}:
            continue  # host diagnostics belong to the controller only
        if control_home.startswith(target + "/"):
            raise DockerError("控制器密钥目录不能与 Worker 共享，请为 DEER_FLOW_WORKER_CONTROL_HOME 配置独立持久卷")
        if mount["Type"] not in {"bind", "volume"}:
            continue
        mounts.append({"Type": mount["Type"], "Source": mount.get("Name") if mount["Type"] == "volume" else mount["Source"], "Target": mount["Destination"], "ReadOnly": not mount.get("RW", True)})
    env = dict(item.split("=", 1) for item in source["Config"].get("Env", []) if "=" in item)
    if not env.get("AUTH_JWT_SECRET"):
        raise DockerError("请在源容器环境变量中设置固定 AUTH_JWT_SECRET，使 Worker 能验证现有登录会话")
    if not mounts:
        raise DockerError("源容器没有共享持久化挂载，无法保证 Worker 访问技能、配置和上传文件")
    if require_shared_archive:
        if env.get("DEER_FLOW_CHAT_ARCHIVE_DIR"):
            archive_path = posixpath.normpath(env["DEER_FLOW_CHAT_ARCHIVE_DIR"])
        elif env.get("DEER_FLOW_HOME"):
            archive_path = posixpath.join(env["DEER_FLOW_HOME"], "data", "chat-archives")
        else:
            project_root = env.get("DEER_FLOW_PROJECT_ROOT") or source["Config"].get("WorkingDir") or settings.working_dir
            archive_path = posixpath.join(project_root, ".deer-flow", "data", "chat-archives")
        if not posixpath.isabs(archive_path):
            workdir = source["Config"].get("WorkingDir") or settings.working_dir
            archive_path = posixpath.join(workdir, archive_path)
        archive_path = posixpath.normpath(archive_path)
        shared = any(
            mount.get("ReadOnly") is False
            and (archive_path == mount["Target"].rstrip("/")
                 or archive_path.startswith(mount["Target"].rstrip("/") + "/"))
            for mount in mounts
        )
        if not shared:
            raise DockerError(
                f"PostgreSQL 聊天归档目录 {archive_path} 必须位于主容器与 Worker 共享的可写持久卷中；"
                "请将 DEER_FLOW_CHAT_ARCHIVE_DIR 指向共享卷目录或挂载 DEER_FLOW_HOME/data"
            )
    for key in list(env):
        if key.startswith("DEER_FLOW_WORKER_") or key in {"DEER_FLOW_REBUILD_WORKFLOW_SCHEMA", "WEB_CONCURRENCY", "UVICORN_WORKERS"}:
            env.pop(key)
    env.update(
        DEER_FLOW_WORKER_DB_POOL_SIZE=str(settings.database_pool_size),
        DEER_FLOW_PG_POOL_MIN="1",
        DEER_FLOW_PG_POOL_MAX=str(settings.postgres_pool_size),
        DEER_FLOW_WORKER_ROLE="online",
        DEER_FLOW_WORKER_TOKEN=token,
        DEER_FLOW_WORKER_CONCURRENCY=str(settings.concurrency),
        DEER_FLOW_WORKER_QUEUE_LIMIT=str(settings.queue_limit),
        DEER_FLOW_WORKER_QUEUE_TIMEOUT=str(settings.queue_timeout),
        SCHEDULED_TASKS_ENABLED="false",
        TASK_ANALYSIS_ENABLED="false",
    )
    return {
        "Image": runtime_image or source["Image"],
        "WorkingDir": settings.working_dir,
        "User": source["Config"].get("User", ""),
        "Entrypoint": [settings.python_executable, "-m", "uvicorn"],
        "Cmd": ["app.gateway.worker_pool.worker:app", "--host", "0.0.0.0", "--port", "8001", "--workers", "1"],
        "Env": [f"{key}={value}" for key, value in env.items()],
        "Labels": {"deerflow.worker-pool": instance, "deerflow.worker-role": "online"},
        "NetworkingConfig": {"EndpointsConfig": {name: {} for name in source["NetworkSettings"]["Networks"] if name not in {"host", "none", "bridge"}}},
        "Healthcheck": {
            "Test": ["CMD", settings.python_executable, "-c", "import os,urllib.request; urllib.request.urlopen(urllib.request.Request('http://127.0.0.1:8001/internal/worker-status',headers={'x-deerflow-worker-token':os.environ['DEER_FLOW_WORKER_TOKEN']}),timeout=4).read()"],
            "Interval": 20_000_000_000, "Timeout": 5_000_000_000, "Retries": 6, "StartPeriod": 90_000_000_000,
        },
        "HostConfig": {
            "Mounts": mounts,
            "NetworkMode": settings.network,
            "Memory": settings.memory_mb * 1024 * 1024,
            "NanoCpus": int(settings.cpu * 1_000_000_000),
            "PidsLimit": 512,
            "RestartPolicy": {"Name": "unless-stopped"},
            "SecurityOpt": ["no-new-privileges:true"],
            "ExtraHosts": source.get("HostConfig", {}).get("ExtraHosts") or [],
            "Dns": source.get("HostConfig", {}).get("Dns") or [],
            "Init": True,
            "LogConfig": {"Type": "json-file", "Config": {"max-size": "10m", "max-file": "3"}},
        },
    }