"""Local persistent worker backend for the sandbox runner runtime."""

from __future__ import annotations

import threading
import time
import venv
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

from mindroom.runtime_env_policy import SANDBOX_RUNTIME_ENV_BY_KEY
from mindroom.tool_system.worker_routing import worker_dir_name
from mindroom.workers.backend import (
    WorkerBackend,
    WorkerBackendError,
    effective_idle_status,
    filter_and_sort_worker_handles,
)
from mindroom.workers.backends._lifecycle import (
    initial_worker_lifecycle_state,
    mark_worker_failed,
    mark_worker_idle,
    mark_worker_ready,
    prepare_worker_ensure_lifecycle,
    read_lifecycle_state,
    touch_worker_lifecycle,
    write_lifecycle_state,
)
from mindroom.workers.backends._metadata_store import (
    list_worker_state_paths,
    load_worker_metadata,
    save_worker_metadata,
)
from mindroom.workers.models import ProgressSink, WorkerHandle, WorkerSpec, WorkerStatus

if TYPE_CHECKING:
    from mindroom.constants import RuntimePaths

_DEFAULT_IDLE_TIMEOUT_SECONDS = 1800.0
_DEFAULT_WORKER_API_ROOT = "/api/sandbox-runner"
_SHARED_INITIALIZATION_LOCK = threading.Lock()
_SHARED_INITIALIZATION_LOCKS: dict[str, threading.Lock] = {}


@dataclass(frozen=True)
class LocalWorkerStatePaths:
    """Filesystem layout for one local worker."""

    root: Path
    workspace: Path
    venv_dir: Path
    cache_dir: Path
    metadata_dir: Path
    metadata_file: Path


@dataclass
class _LocalWorkerMetadata:
    worker_id: str
    worker_key: str
    endpoint: str
    backend_name: str
    created_at: float
    last_used_at: float
    status: WorkerStatus
    last_started_at: float | None = None
    startup_count: int = 0
    failure_count: int = 0
    failure_reason: str | None = None


def _default_worker_root(runtime_paths: RuntimePaths) -> Path:
    return runtime_paths.storage_root.resolve() / "workers"


def _read_idle_timeout_seconds(runtime_paths: RuntimePaths) -> float:
    raw_timeout = runtime_paths.env_value(
        SANDBOX_RUNTIME_ENV_BY_KEY["worker_idle_timeout_seconds"],
        default=str(_DEFAULT_IDLE_TIMEOUT_SECONDS),
    ) or str(_DEFAULT_IDLE_TIMEOUT_SECONDS)
    try:
        timeout = float(raw_timeout)
    except ValueError:
        timeout = _DEFAULT_IDLE_TIMEOUT_SECONDS
    return max(1.0, timeout)


def _normalize_worker_api_root(raw_endpoint: str) -> str:
    normalized = raw_endpoint.strip() or _DEFAULT_WORKER_API_ROOT
    normalized = normalized.rstrip("/")
    if normalized.endswith("/execute"):
        normalized = normalized.removesuffix("/execute")
    return normalized or _DEFAULT_WORKER_API_ROOT


def _read_worker_api_root(runtime_paths: RuntimePaths) -> str:
    raw_api_root = runtime_paths.env_value(
        SANDBOX_RUNTIME_ENV_BY_KEY["worker_endpoint"],
        default=_DEFAULT_WORKER_API_ROOT,
    )
    return _normalize_worker_api_root(raw_api_root or _DEFAULT_WORKER_API_ROOT)


def local_worker_state_paths_for_root(state_root: Path) -> LocalWorkerStatePaths:
    """Return the filesystem paths owned by one concrete worker runtime root."""
    resolved_root = state_root.expanduser().resolve()
    metadata_dir = resolved_root / "metadata"
    return LocalWorkerStatePaths(
        root=resolved_root,
        workspace=resolved_root / "workspace",
        venv_dir=resolved_root / "venv",
        cache_dir=resolved_root / "cache",
        metadata_dir=metadata_dir,
        metadata_file=metadata_dir / "worker.json",
    )


def _local_worker_state_paths(worker_key: str, *, worker_root: Path) -> LocalWorkerStatePaths:
    """Return the runtime-local filesystem paths owned by one worker key."""
    resolved_root = worker_root.expanduser().resolve()
    return local_worker_state_paths_for_root(resolved_root / worker_dir_name(worker_key))


def local_worker_state_paths_from_handle(handle: WorkerHandle) -> LocalWorkerStatePaths:
    """Resolve local runtime paths from a local worker handle."""
    state_root = handle.debug_metadata.get("state_root")
    if state_root is None:
        msg = f"Worker '{handle.worker_key}' does not expose local state metadata."
        raise WorkerBackendError(msg)
    return local_worker_state_paths_for_root(Path(state_root))


def _ensure_local_worker_state(paths: LocalWorkerStatePaths) -> None:
    """Create the persistent directories and venv for one worker runtime root."""
    paths.workspace.mkdir(parents=True, exist_ok=True)
    paths.cache_dir.mkdir(parents=True, exist_ok=True)
    paths.metadata_dir.mkdir(parents=True, exist_ok=True)
    if (paths.venv_dir / "bin" / "python").exists():
        return

    builder = venv.EnvBuilder(with_pip=True, system_site_packages=True)
    builder.create(paths.venv_dir)


def ensure_local_worker_state_locked(paths: LocalWorkerStatePaths) -> None:
    """Create one worker runtime root under a shared per-worker initialization lock."""
    with _shared_worker_initialization_lock(paths):
        _ensure_local_worker_state(paths)


def _shared_worker_initialization_lock(paths: LocalWorkerStatePaths) -> threading.Lock:
    lock_key = str(paths.root)
    with _SHARED_INITIALIZATION_LOCK:
        worker_lock = _SHARED_INITIALIZATION_LOCKS.get(lock_key)
        if worker_lock is None:
            worker_lock = threading.Lock()
            _SHARED_INITIALIZATION_LOCKS[lock_key] = worker_lock
    return worker_lock


class _LocalWorkerBackend:
    """Persistent local worker backend used by the sandbox runner."""

    backend_name = "local_sandbox_runner"

    def __init__(
        self,
        *,
        worker_root: Path,
        api_root: str,
        idle_timeout_seconds: float,
    ) -> None:
        self.worker_root = worker_root.expanduser().resolve()
        self.api_root = _normalize_worker_api_root(api_root)
        self.idle_timeout_seconds = max(1.0, idle_timeout_seconds)
        self.worker_root.mkdir(parents=True, exist_ok=True)
        self._lock = threading.Lock()

    def shutdown(self) -> None:
        """Local worker state is persistent; manager replacement does not need extra teardown."""
        return

    def ensure_worker(
        self,
        spec: WorkerSpec,
        *,
        now: float | None = None,
        progress_sink: ProgressSink | None = None,
    ) -> WorkerHandle:
        """Resolve or create one local worker."""
        del progress_sink
        timestamp = time.time() if now is None else now
        paths = _local_worker_state_paths(spec.worker_key, worker_root=self.worker_root)
        worker_lock = self._worker_lock(paths)

        with worker_lock:
            with self._lock:
                metadata = self._load_metadata(paths) or self._default_metadata(spec.worker_key, timestamp)
                should_restart = self._effective_status(metadata, timestamp) != "ready"
                write_lifecycle_state(
                    metadata,
                    prepare_worker_ensure_lifecycle(
                        read_lifecycle_state(metadata),
                        now=timestamp,
                        should_restart=should_restart,
                    ),
                )
                self._save_metadata(paths, metadata)

            try:
                self._ensure_worker_state(paths)
            except Exception as exc:
                failure_reason = f"Failed to initialize worker '{spec.worker_key}': {exc}"
                with self._lock:
                    self._record_failure_locked(paths, spec.worker_key, failure_reason, now=timestamp)
                raise WorkerBackendError(failure_reason) from exc

            with self._lock:
                write_lifecycle_state(metadata, mark_worker_ready(read_lifecycle_state(metadata), now=timestamp))
                self._save_metadata(paths, metadata)
                return self._to_handle(metadata, paths, now=timestamp)

    def touch_worker(self, worker_key: str, *, now: float | None = None) -> WorkerHandle | None:
        """Refresh last-used bookkeeping for one local worker."""
        timestamp = time.time() if now is None else now
        paths = _local_worker_state_paths(worker_key, worker_root=self.worker_root)
        worker_lock = self._worker_lock(paths)
        with worker_lock, self._lock:
            metadata = self._load_metadata(paths)
            if metadata is None:
                return None
            write_lifecycle_state(metadata, touch_worker_lifecycle(read_lifecycle_state(metadata), now=timestamp))
            self._save_metadata(paths, metadata)
            return self._to_handle(metadata, paths, now=timestamp)

    def list_workers(self, *, include_idle: bool = True, now: float | None = None) -> list[WorkerHandle]:
        """List known local workers."""
        timestamp = time.time() if now is None else now
        with self._lock:
            handles = [
                self._to_handle(metadata, paths, now=timestamp)
                for paths in self._metadata_paths()
                if (metadata := self._load_metadata(paths)) is not None
            ]

        return filter_and_sort_worker_handles(handles, include_idle)

    def cleanup_idle_workers(self, *, now: float | None = None) -> list[WorkerHandle]:
        """Mark timed-out local workers idle."""
        timestamp = time.time() if now is None else now
        cleaned_workers: list[WorkerHandle] = []

        with self._lock:
            for paths in self._metadata_paths():
                metadata = self._load_metadata(paths)
                if metadata is None:
                    continue
                if metadata.status == "ready" and self._effective_status(metadata, timestamp) == "idle":
                    write_lifecycle_state(metadata, mark_worker_idle(read_lifecycle_state(metadata)))
                    self._save_metadata(paths, metadata)
                    cleaned_workers.append(self._to_handle(metadata, paths, now=timestamp))

        return filter_and_sort_worker_handles(cleaned_workers, True)

    def record_failure(self, worker_key: str, failure_reason: str, *, now: float | None = None) -> WorkerHandle:
        """Persist one local worker failure."""
        timestamp = time.time() if now is None else now
        paths = _local_worker_state_paths(worker_key, worker_root=self.worker_root)
        worker_lock = self._worker_lock(paths)

        with worker_lock, self._lock:
            return self._record_failure_locked(paths, worker_key, failure_reason, now=timestamp)

    def _worker_lock(self, paths: LocalWorkerStatePaths) -> threading.Lock:
        return _shared_worker_initialization_lock(paths)

    def _default_metadata(self, worker_key: str, now: float) -> _LocalWorkerMetadata:
        lifecycle = initial_worker_lifecycle_state(now=now)
        return _LocalWorkerMetadata(
            worker_id=worker_dir_name(worker_key),
            worker_key=worker_key,
            endpoint=f"{self.api_root}/execute",
            backend_name=self.backend_name,
            created_at=lifecycle.created_at,
            last_used_at=lifecycle.last_used_at,
            status=lifecycle.status,
        )

    def _ensure_worker_state(self, paths: LocalWorkerStatePaths) -> None:
        _ensure_local_worker_state(paths)

    def _metadata_paths(self) -> list[LocalWorkerStatePaths]:
        return list_worker_state_paths(
            self.worker_root,
            state_paths_from_root=local_worker_state_paths_for_root,
        )

    def _load_metadata(self, paths: LocalWorkerStatePaths) -> _LocalWorkerMetadata | None:
        return load_worker_metadata(paths, metadata_type=_LocalWorkerMetadata)

    def _save_metadata(self, paths: LocalWorkerStatePaths, metadata: _LocalWorkerMetadata) -> None:
        save_worker_metadata(paths, metadata)

    def _effective_status(self, metadata: _LocalWorkerMetadata, now: float) -> WorkerStatus:
        return effective_idle_status(metadata.status, metadata.last_used_at, self.idle_timeout_seconds, now)

    def _record_failure_locked(
        self,
        paths: LocalWorkerStatePaths,
        worker_key: str,
        failure_reason: str,
        *,
        now: float,
    ) -> WorkerHandle:
        metadata = self._load_metadata(paths) or self._default_metadata(worker_key, now)
        write_lifecycle_state(
            metadata,
            mark_worker_failed(read_lifecycle_state(metadata), now=now, failure_reason=failure_reason),
        )
        self._save_metadata(paths, metadata)
        return self._to_handle(metadata, paths, now=now)

    def _to_handle(self, metadata: _LocalWorkerMetadata, paths: LocalWorkerStatePaths, *, now: float) -> WorkerHandle:
        return WorkerHandle(
            worker_id=metadata.worker_id,
            worker_key=metadata.worker_key,
            endpoint=metadata.endpoint,
            auth_token=None,
            status=self._effective_status(metadata, now),
            backend_name=metadata.backend_name,
            last_used_at=metadata.last_used_at,
            created_at=metadata.created_at,
            last_started_at=metadata.last_started_at,
            expires_at=None,
            startup_count=metadata.startup_count,
            failure_count=metadata.failure_count,
            failure_reason=metadata.failure_reason,
            debug_metadata={
                "api_root": self.api_root,
                "state_root": str(paths.root),
            },
        )


_local_worker_manager: WorkerBackend | None = None
_local_worker_manager_config: tuple[str, str, float] | None = None
_local_worker_manager_lock = threading.Lock()


def get_local_worker_manager(runtime_paths: RuntimePaths) -> WorkerBackend:
    """Return the local sandbox worker manager for the current config."""
    global _local_worker_manager, _local_worker_manager_config

    worker_root = _default_worker_root(runtime_paths)
    api_root = _read_worker_api_root(runtime_paths)
    idle_timeout_seconds = _read_idle_timeout_seconds(runtime_paths)
    config = (str(worker_root), api_root, idle_timeout_seconds)

    with _local_worker_manager_lock:
        if _local_worker_manager is None or _local_worker_manager_config != config:
            _local_worker_manager = _LocalWorkerBackend(
                worker_root=worker_root,
                api_root=api_root,
                idle_timeout_seconds=idle_timeout_seconds,
            )
            _local_worker_manager_config = config

    return _local_worker_manager
