"""Workspace resolution and scaffolding helpers for agents."""

from __future__ import annotations

import shutil
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

from mindroom.constants import RuntimePaths, config_relative_path
from mindroom.runtime_env_policy import SANDBOX_RUNTIME_ENV_BY_KEY

if TYPE_CHECKING:
    from collections.abc import Collection, Mapping

    from mindroom.config.main import Config

_MIND_TEMPLATE_DIR = Path(__file__).resolve().parent / "cli" / "templates" / "mind_data"

# Agent builds now run on worker threads (#1260), so concurrent scaffolding of
# the same workspace must not interleave template copies or link reconciliation.
_WORKSPACE_MUTATION_LOCK = threading.Lock()


@dataclass(frozen=True)
class ResolvedAgentWorkspace:
    """Resolved workspace paths for one agent in one execution scope."""

    root: Path
    context_files: tuple[Path, ...]
    file_memory_path: Path | None


@dataclass(frozen=True)
class _EffectiveAgentWorkspace:
    root_path: str
    template_dir: Path | None
    context_files: tuple[str, ...]
    file_memory_path: str | None


def validate_local_copy_source_path(
    source_path: Path,
    *,
    field_name: str,
) -> Path:
    """Resolve one existing local copy-source path and reject symlink traversal."""
    lexical_source_path = source_path.expanduser()
    absolute_source_path = (
        lexical_source_path if lexical_source_path.is_absolute() else (Path.cwd() / lexical_source_path)
    )
    current = Path(absolute_source_path.anchor) if absolute_source_path.anchor else Path()
    for part in absolute_source_path.parts[1:] if absolute_source_path.anchor else absolute_source_path.parts:
        current = current / part
        if current.is_symlink():
            msg = f"{field_name} must not contain symlinks: {current}"
            raise ValueError(msg)
    resolved_source_path = absolute_source_path.resolve()
    if not resolved_source_path.exists():
        msg = f"{field_name} does not exist: {resolved_source_path}"
        raise ValueError(msg)
    return resolved_source_path


def iter_local_copy_source_entries(source_dir: Path) -> list[tuple[Path, Path]]:
    """Return copy-source entries in deterministic order without following symlinks."""
    entries: list[tuple[Path, Path]] = []

    def _walk(current_dir: Path) -> None:
        for source_path in sorted(current_dir.iterdir()):
            relative_path = source_path.relative_to(source_dir)
            entries.append((source_path, relative_path))
            if source_path.is_dir() and not source_path.is_symlink():
                _walk(source_path)

    _walk(source_dir)
    return entries


def validate_local_copy_source_dir(
    source_dir: Path,
    *,
    field_name: str,
) -> Path:
    """Resolve one local copy-source directory and reject symlinked content."""
    resolved_source_dir = validate_local_copy_source_path(
        source_dir,
        field_name=field_name,
    )
    if not resolved_source_dir.is_dir():
        msg = f"{field_name} is not a directory: {resolved_source_dir}"
        raise ValueError(msg)
    for source_path, _ in iter_local_copy_source_entries(resolved_source_dir):
        if source_path.is_symlink():
            msg = f"{field_name} must not contain symlinks: {source_path}"
            raise ValueError(msg)
    return resolved_source_dir


def resolve_relative_path_within_root(
    root: Path,
    relative_path: str | Path,
    *,
    field_name: str,
    root_label: str = "canonical root",
) -> Path:
    """Resolve one relative path under a canonical root and reject symlink escapes."""
    lexical_root = root.expanduser()
    resolved_root = lexical_root.resolve()
    candidate_path = lexical_root / relative_path
    current = lexical_root
    for part in Path(relative_path).parts:
        current = current / part
        if current.is_symlink():
            msg = f"{field_name} must stay within the {root_label}: {resolved_root}"
            raise ValueError(msg)
    candidate = candidate_path.resolve()
    if not candidate.is_relative_to(resolved_root):
        msg = f"{field_name} must stay within the {root_label}: {resolved_root}"
        raise ValueError(msg)
    return candidate


def resolve_relative_path_within_root_preserving_leaf(
    root: Path,
    relative_path: str | Path,
    *,
    field_name: str,
    root_label: str = "canonical root",
) -> Path:
    """Resolve one relative path under a canonical root without following the final component."""
    lexical_root = root.expanduser()
    resolved_root = lexical_root.resolve()
    candidate_path = lexical_root / relative_path
    relative = Path(relative_path)
    if relative.is_absolute():
        msg = f"{field_name} must stay within the {root_label}: {resolved_root}"
        raise ValueError(msg)
    if relative == Path():
        return resolved_root

    current = lexical_root
    for index, part in enumerate(relative.parts):
        if part == "..":
            msg = f"{field_name} must stay within the {root_label}: {resolved_root}"
            raise ValueError(msg)
        current = current / part
        if index < len(relative.parts) - 1 and current.is_symlink():
            msg = f"{field_name} must stay within the {root_label}: {resolved_root}"
            raise ValueError(msg)

    candidate_parent = candidate_path.parent.resolve()
    if not candidate_parent.is_relative_to(resolved_root):
        msg = f"{field_name} must stay within the {root_label}: {resolved_root}"
        raise ValueError(msg)
    return candidate_path


def resolve_workspace_relative_path(
    root: Path,
    relative_path: str | Path,
    *,
    field_name: str,
) -> Path:
    """Resolve one workspace-relative path and reject symlink escapes."""
    return resolve_relative_path_within_root(
        root,
        relative_path,
        field_name=field_name,
        root_label="workspace root",
    )


def validate_workspace_template_dir(template_dir: Path) -> Path:
    """Resolve one workspace template directory and ensure it exists."""
    return validate_local_copy_source_dir(
        template_dir,
        field_name="Workspace template directory",
    )


def _iter_workspace_template_entries(template_dir: Path) -> list[tuple[Path, Path]]:
    """Return template entries in deterministic order without following symlinks."""
    return iter_local_copy_source_entries(template_dir)


def _copy_workspace_template(
    workspace_path: Path,
    *,
    template_dir: Path,
    force: bool = False,
) -> None:
    """Copy a local template directory into a workspace root."""
    workspace_path.mkdir(parents=True, exist_ok=True)
    resolved_template_dir = validate_workspace_template_dir(template_dir)

    with _WORKSPACE_MUTATION_LOCK:
        for source_path, relative_path in _iter_workspace_template_entries(resolved_template_dir):
            destination_path = resolve_relative_path_within_root(
                workspace_path,
                relative_path,
                field_name="workspace template destination",
                root_label="workspace root",
            )
            if source_path.is_dir():
                destination_path.mkdir(parents=True, exist_ok=True)
                continue
            if destination_path.exists() and not force:
                continue
            destination_path.parent.mkdir(parents=True, exist_ok=True)
            # Publish atomically so concurrent readers never see partial files.
            temp_path = destination_path.with_name(f".{destination_path.name}.tmp")
            try:
                shutil.copy2(source_path, temp_path)
                temp_path.replace(destination_path)
            finally:
                temp_path.unlink(missing_ok=True)


def ensure_workspace_template(
    workspace_path: Path,
    *,
    template: str,
    force: bool = False,
) -> None:
    """Create the built-in Mind workspace template used by config init."""
    if template != "mind":
        msg = f"Unsupported workspace template: {template}"
        raise ValueError(msg)
    _copy_workspace_template(workspace_path, template_dir=_MIND_TEMPLATE_DIR, force=force)
    (workspace_path / "memory").mkdir(parents=True, exist_ok=True)


def _build_workspace_knowledge_links(
    *,
    workspace_root: Path,
    knowledge_root: Path,
    knowledge_paths: Mapping[str, Path],
) -> dict[Path, Path]:
    desired_links: dict[Path, Path] = {}
    for base_id, target_path in knowledge_paths.items():
        link_path = resolve_relative_path_within_root_preserving_leaf(
            knowledge_root,
            base_id,
            field_name="workspace knowledge link",
        )
        resolved_target = target_path.expanduser().resolve()
        if not resolved_target.is_relative_to(workspace_root):
            continue
        if resolved_target == link_path or resolved_target.is_relative_to(link_path):
            continue
        if link_path.is_relative_to(resolved_target):
            continue
        desired_links[link_path] = resolved_target
    return desired_links


def _remove_stale_workspace_knowledge_links(
    knowledge_root: Path,
    *,
    protected_paths: Collection[Path],
    workspace_root: Path,
    desired_links: Mapping[Path, Path],
) -> None:
    for existing_path in knowledge_root.iterdir():
        if (
            existing_path.is_symlink()
            and existing_path not in desired_links
            and existing_path not in protected_paths
            and existing_path.resolve().is_relative_to(workspace_root)
        ):
            existing_path.unlink()


def _apply_workspace_knowledge_links(desired_links: Mapping[Path, Path]) -> None:
    for link_path, resolved_target in desired_links.items():
        if link_path.is_symlink():
            if link_path.resolve() == resolved_target:
                continue
            link_path.unlink()
        elif link_path.exists():
            if link_path.resolve() == resolved_target:
                continue
            msg = f"Workspace knowledge link path already exists and is not a symlink: {link_path}"
            raise ValueError(msg)
        if resolved_target == link_path:
            continue
        link_path.parent.mkdir(parents=True, exist_ok=True)
        link_path.symlink_to(resolved_target, target_is_directory=True)


def ensure_workspace_knowledge_links(
    workspace_path: Path,
    *,
    knowledge_paths: Mapping[str, Path],
    protected_paths: Collection[Path] = (),
) -> None:
    """Expose canonical workspace-local paths for bound knowledge rooted inside the workspace.

    Each knowledge base becomes visible under ``<workspace>/knowledge/<base_id>``.
    Targets outside the workspace are intentionally excluded because the default
    file-aware tools enforce workspace containment after resolving symlinks.
    """
    workspace_path.mkdir(parents=True, exist_ok=True)
    workspace_root = workspace_path.resolve()
    knowledge_root = resolve_workspace_relative_path(
        workspace_path,
        "knowledge",
        field_name="workspace knowledge root",
    )
    if not knowledge_paths and not knowledge_root.exists():
        return
    knowledge_root.mkdir(parents=True, exist_ok=True)
    with _WORKSPACE_MUTATION_LOCK:
        desired_links = _build_workspace_knowledge_links(
            workspace_root=workspace_root,
            knowledge_root=knowledge_root,
            knowledge_paths=knowledge_paths,
        )
        _remove_stale_workspace_knowledge_links(
            knowledge_root,
            protected_paths=protected_paths,
            workspace_root=workspace_root,
            desired_links=desired_links,
        )
        _apply_workspace_knowledge_links(desired_links)


def _private_root_name(agent_name: str, config: Config) -> str:
    agent_config = config.agents.get(agent_name)
    if agent_config is None or agent_config.private is None or agent_config.private.root is None:
        return f"{agent_name}_data"
    return agent_config.private.root


def _effective_workspace(
    agent_name: str,
    config: Config,
    *,
    runtime_paths: RuntimePaths,
) -> _EffectiveAgentWorkspace | None:
    agent_config = config.agents.get(agent_name)
    if agent_config is None or agent_config.private is None:
        return None
    private_config = agent_config.private
    return _EffectiveAgentWorkspace(
        root_path=_private_root_name(agent_name, config),
        template_dir=(
            config_relative_path(private_config.template_dir, runtime_paths)
            if private_config.template_dir is not None
            else None
        ),
        context_files=tuple(private_config.context_files or ()),
        file_memory_path="." if config.resolve_entity(agent_name).memory_backend == "file" else None,
    )


def _template_unavailable_for_dedicated_worker(template_dir: Path, runtime_paths: RuntimePaths) -> bool:
    """Return whether a dedicated worker should rely on the control plane's existing scaffold."""
    return (
        runtime_paths.env_flag(SANDBOX_RUNTIME_ENV_BY_KEY["runner_mode"])
        and bool(runtime_paths.env_value(SANDBOX_RUNTIME_ENV_BY_KEY["dedicated_worker_key"], default=""))
        and not template_dir.expanduser().is_dir()
    )


def _resolve_workspace(
    agent_name: str,
    config: Config,
    *,
    runtime_paths: RuntimePaths,
    state_storage_path: Path,
    use_state_storage_path: bool,
    create: bool,
) -> ResolvedAgentWorkspace | None:
    agent_config = config.agents.get(agent_name)
    if agent_config is None:
        return None

    if agent_config.private is None:
        if config.resolve_entity(agent_name).memory_backend != "file":
            return None
        root = resolve_workspace_relative_path(
            state_storage_path,
            "workspace",
            field_name="agent workspace root",
        )
        if create:
            root.mkdir(parents=True, exist_ok=True)
        return ResolvedAgentWorkspace(
            root=root,
            context_files=(),
            file_memory_path=root,
        )

    workspace = _effective_workspace(agent_name, config, runtime_paths=runtime_paths)
    assert workspace is not None

    if not use_state_storage_path:
        msg = f"Private agent '{agent_name}' requires an active execution identity to resolve requester-local state"
        raise ValueError(msg)

    root = resolve_workspace_relative_path(
        state_storage_path,
        workspace.root_path,
        field_name="private.root",
    )
    template_dir = workspace.template_dir
    if create:
        root.mkdir(parents=True, exist_ok=True)
        if template_dir is not None:
            assert template_dir is not None
            if not _template_unavailable_for_dedicated_worker(template_dir, runtime_paths):
                _copy_workspace_template(root, template_dir=template_dir)

    context_files = tuple(
        resolve_workspace_relative_path(
            root,
            relative_path,
            field_name="private.context_files",
        )
        for relative_path in workspace.context_files
    )
    file_memory_path = (
        resolve_workspace_relative_path(
            root,
            workspace.file_memory_path,
            field_name="private file memory path",
        )
        if workspace.file_memory_path is not None
        else None
    )
    if create and file_memory_path is not None:
        file_memory_path.mkdir(parents=True, exist_ok=True)

    return ResolvedAgentWorkspace(
        root=root,
        context_files=context_files,
        file_memory_path=file_memory_path,
    )


def resolve_agent_workspace_from_state_path(
    agent_name: str,
    config: Config,
    *,
    runtime_paths: RuntimePaths,
    state_storage_path: Path,
    use_state_storage_path: bool,
    create: bool = False,
) -> ResolvedAgentWorkspace | None:
    """Resolve one agent workspace when the caller already knows the state root."""
    return _resolve_workspace(
        agent_name,
        config,
        runtime_paths=runtime_paths,
        state_storage_path=state_storage_path.expanduser(),
        use_state_storage_path=use_state_storage_path,
        create=create,
    )
