headquarter/apps/api/src/services/shared/workspace_manager.py

"""Workspace lifecycle management service."""

from __future__ import annotations

import contextlib
import logging
import os
import shutil
import stat
import uuid
from dataclasses import dataclass
from datetime import datetime
from typing import TYPE_CHECKING

from sqlalchemy import select

from src.models import Workspace
from src.services.git.git_service import GitService
from src.services.shared.ssh_keys import _get_fernet
from src.utils.git_url_parser import extract_base_repo_url

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession

    from src.models import GitRepository
    from src.models import ToolInstance

logger = logging.getLogger(__name__)


@dataclass
class SyncResult:
    """Result of a workspace sync operation."""

    branch_deleted: bool = False


class WorkspaceHasInstancesError(Exception):
    """Raised when attempting to delete a workspace with running instances."""

    def __init__(self, instances: list[dict]) -> None:
        self.instances = instances
        super().__init__(f"Workspace has {len(instances)} running tool instance(s)")


class WorkspaceManager:
    """Manages workspace lifecycle: create, delete, sync, validate."""

    BASE_PATH = "/data/working-copies"

    @staticmethod
    def _repo_directory_name(repo: "GitRepository") -> str:
        """Return the directory name git would create for a standard clone.

        Prefers the name parsed from the remote URL and falls back to the
        user-provided repository name when no remote URL is available.
        """
        if repo.remote_url:
            base_url = extract_base_repo_url(repo.remote_url) or repo.remote_url
            name = base_url.rstrip("/").split("/")[-1]
            if name.endswith(".git"):
                name = name[:-4]
            if name:
                return name
        return repo.name

    def _workspace_path(self, workspace_id: uuid.UUID, repo: "GitRepository") -> str:
        """Return the filesystem path for a workspace.

        Layout: /data/working-copies/{workspace_id}/{repo_name}/
        The workspace_id prevents collisions between workspaces, and the
        repo_name matches the directory git clone naturally creates.
        """
        return os.path.join(
            self.BASE_PATH, str(workspace_id), self._repo_directory_name(repo)
        )

    async def create(
        self,
        repo: GitRepository,
        user_id: uuid.UUID,
        name: str,
        branch: str = "main",
        session: AsyncSession | None = None,
    ) -> Workspace:
        """Clone repo to workspace path and create DB record.

        Args:
            repo: The git repository to clone.
            user_id: The owner user ID.
            name: The workspace name (unique per repo).
            branch: The branch to clone (default: "main").
            session: Database session for loading SSH keys.

        Returns:
            The created Workspace record.

        Raises:
            RuntimeError: If git clone fails.
        """
        if not repo.remote_url:
            raise ValueError("Repository has no remote URL")

        repo_dir_name = self._repo_directory_name(repo)
        workspace_id = uuid.uuid4()
        parent_path = os.path.join(self.BASE_PATH, str(workspace_id))
        clone_target = parent_path  # git clone creates {repo_dir_name}/ inside this

        logger.info(
            "Creating workspace: id=%s name=%s repo=%s branch=%s",
            workspace_id,
            name,
            repo.id,
            branch,
        )

        # Remove stale directory from previous failed/aborted clone
        if os.path.exists(parent_path):
            logger.warning("Removing stale workspace directory: %s", parent_path)
            shutil.rmtree(parent_path, ignore_errors=True)

        os.makedirs(parent_path, exist_ok=True)
        # Ensure container users (various UIDs) can write to workspace dirs
        with contextlib.suppress(OSError):
            os.chmod(parent_path, 0o777)

        # Load SSH key if repo has one
        ssh_key = None
        if getattr(repo, "ssh_key_id", None) and session is not None:
            from src.models import SSHKey

            result = await session.execute(
                select(SSHKey).where(SSHKey.id == repo.ssh_key_id)
            )
            ssh_key_obj = result.scalar_one_or_none()
            if ssh_key_obj:
                fernet = _get_fernet()
                ssh_key = fernet.decrypt(
                    ssh_key_obj.private_key_encrypted.encode()
                ).decode()

        await GitService.clone(repo.remote_url, branch, clone_target, ssh_key=ssh_key)

        expected_path = os.path.join(parent_path, repo_dir_name)
        if not os.path.isdir(expected_path):
            # git clone can create a different directory name than expected for
            # some URL shapes; fall back to the single directory git created.
            entries = [
                entry
                for entry in os.listdir(parent_path)
                if os.path.isdir(os.path.join(parent_path, entry))
            ]
            if len(entries) == 1:
                expected_path = os.path.join(parent_path, entries[0])
            else:
                raise RuntimeError(
                    f"Expected git clone to create {repo_dir_name}/ under "
                    f"{parent_path}, but found: {entries}"
                )

        self._make_world_writable(expected_path)

        workspace = Workspace(
            id=workspace_id,
            name=name,
            repo_id=repo.id,
            user_id=user_id,
            branch=branch,
            path=expected_path,
            status="ready",
            last_sync_at=datetime.now(),
        )
        logger.info("Workspace created: %s at %s", workspace.id, workspace.path)
        return workspace

    async def delete(
        self,
        workspace: Workspace,
        force: bool = False,
        session: AsyncSession | None = None,
    ) -> None:
        """Delete a workspace and all associated tool instances.

        Args:
            workspace: The workspace to delete.
            force: If True, delete even if instances exist.
            session: The database session (required for checking instances).

        Raises:
            WorkspaceHasInstancesError: If instances exist and force=False.
        """
        if session is None:
            raise ValueError("session is required for delete")

        instances = await self._get_instances(workspace, session)
        if instances and not force:
            raise WorkspaceHasInstancesError(
                [{"id": str(i.id), "name": i.name} for i in instances]
            )

        # Stop and delete all instances
        for instance in instances:
            await self._stop_and_delete_instance(instance, session)

        # Delete directory
        if os.path.exists(workspace.path):
            shutil.rmtree(workspace.path, ignore_errors=True)
            logger.info("Deleted workspace directory: %s", workspace.path)

        # Delete record
        await session.delete(workspace)
        logger.info("Deleted workspace record: %s", workspace.id)

    async def sync(
        self, workspace: Workspace, session: AsyncSession | None = None
    ) -> SyncResult:
        """Sync a workspace with its remote.

        Args:
            workspace: The workspace to sync.
            session: Database session for loading SSH keys.

        Returns:
            SyncResult indicating whether the branch was deleted.

        Raises:
            RuntimeError: If git operations fail.
        """
        logger.info("Syncing workspace: %s", workspace.id)

        # Load SSH key if repo has one
        ssh_key = None
        if session is not None:
            from src.models import GitRepository
            from src.models import SSHKey

            repo = await session.get(GitRepository, workspace.repo_id)
            if repo and getattr(repo, "ssh_key_id", None):
                result = await session.execute(
                    select(SSHKey).where(SSHKey.id == repo.ssh_key_id)
                )
                ssh_key_obj = result.scalar_one_or_none()
                if ssh_key_obj:
                    fernet = _get_fernet()
                    ssh_key = fernet.decrypt(
                        ssh_key_obj.private_key_encrypted.encode()
                    ).decode()

        await GitService.fetch(workspace.path, ssh_key=ssh_key)

        if not GitService.branch_exists_remotely(
            workspace.path, workspace.branch, ssh_key=ssh_key
        ):
            return SyncResult(branch_deleted=True)

        await GitService.pull(workspace.path, workspace.branch, ssh_key=ssh_key)
        self._make_world_writable(workspace.path)

        workspace.last_sync_at = datetime.now()
        logger.info("Workspace synced: %s", workspace.id)
        return SyncResult(branch_deleted=False)

    def _make_world_writable(self, path: str) -> None:
        """Recursively make path readable/writable/traversable by any UID.

        Directories get 777 (traversable). Files get rw for all while
        preserving any existing execute bits.
        """
        with contextlib.suppress(OSError):
            os.chmod(path, 0o777)
        for root, dirs, files in os.walk(path):
            for d in dirs:
                dpath = os.path.join(root, d)
                with contextlib.suppress(OSError):
                    os.chmod(dpath, 0o777)
            for f in files:
                fpath = os.path.join(root, f)
                with contextlib.suppress(OSError):
                    mode = os.stat(fpath).st_mode
                    # Preserve execute bits, ensure read+write for all
                    new_mode = (mode & stat.S_IXUSR) | 0o666
                    if mode & stat.S_IXGRP:
                        new_mode |= stat.S_IXGRP
                    if mode & stat.S_IXOTH:
                        new_mode |= stat.S_IXOTH
                    os.chmod(fpath, new_mode)

    async def _get_instances(
        self,
        workspace: Workspace,
        session: AsyncSession,
    ) -> list[ToolInstance]:
        """Get all tool instances associated with this workspace."""
        from src.models import ToolInstance

        result = await session.execute(
            select(ToolInstance).where(ToolInstance.workspace_id == workspace.id)
        )
        return list(result.scalars().all())

    async def _stop_and_delete_instance(
        self,
        instance: ToolInstance,
        session: AsyncSession,
    ) -> None:
        """Stop and delete a tool instance.

        Delegates to the instance lifecycle service to ensure containers,
        tunnels, and on-disk files are cleaned up.
        """
        # Local import avoids a circular dependency between workspace and
        # instance lifecycle modules.
        from src.services.tool.instance_service import delete_tool_instance

        try:
            await delete_tool_instance(
                session=session,
                user_id=instance.owner_id,
                project_id=instance.project_id,
                repo_id=instance.repository_id,
                instance_id=instance.id,
                force=True,
            )
            logger.info("Stopped and deleted instance %s", instance.id)
        except Exception as exc:
            logger.error(
                "Failed to stop/delete instance %s during workspace cleanup: %s",
                instance.id,
                exc,
            )

    async def ensure_instance_workspace(
        self,
        instance: ToolInstance,
        session: AsyncSession,
    ) -> Workspace:
        """Return the workspace for an instance, creating/binding one if needed.

        This migrates legacy instances that were created before workspaces
        existed. Clone-mode instances have their existing clone moved into a
        workspace, while mount-mode instances get a fresh workspace from the
        canonical repository.

        The operation is best-effort: failures are re-raised so the caller can
        decide whether to continue with a legacy fallback path.
        """
        if instance.workspace_id is not None:
            workspace = await session.get(Workspace, instance.workspace_id)
            if workspace is not None:
                return workspace

        from src.models import GitRepository

        repo = await session.get(GitRepository, instance.repository_id)
        if repo is None:
            raise RuntimeError(f"Repository {instance.repository_id} not found")

        base_name = (
            f"{instance.name}-migrated"
            if instance.clone_mode == "clone"
            else f"{instance.name}-legacy"
        )
        name = base_name
        counter = 1
        while await self._workspace_name_exists(session, repo.id, name):
            name = f"{base_name}-{counter}"
            counter += 1

        if instance.clone_mode == "clone":
            workspace = await self._migrate_clone_into_workspace(
                instance, repo, session, name
            )
        else:
            workspace = await self.create(
                repo=repo,
                user_id=instance.owner_id,
                name=name,
                branch=instance.branch or "main",
                session=session,
            )

        instance.workspace_id = workspace.id
        instance.clone_mode = None
        session.add(instance)
        await session.commit()
        await session.refresh(instance)
        logger.info(
            "Migrated instance %s to workspace %s (%s)",
            instance.id,
            workspace.id,
            name,
        )
        return workspace

    async def _workspace_name_exists(
        self,
        session: AsyncSession,
        repo_id: uuid.UUID,
        name: str,
    ) -> bool:
        """Check whether a workspace name already exists for a repository."""
        result = await session.execute(
            select(Workspace).where(
                Workspace.repo_id == repo_id,
                Workspace.name == name,
            )
        )
        return result.scalar_one_or_none() is not None

    async def _migrate_clone_into_workspace(
        self,
        instance: ToolInstance,
        repo: "GitRepository",
        session: AsyncSession,
        name: str,
    ) -> Workspace:
        """Move an existing clone-mode repo into a new workspace path."""
        import shutil

        if not instance.compose_path:
            raise RuntimeError("Instance has no compose path")

        instance_dir = os.path.dirname(instance.compose_path)
        clone_path = os.path.join(instance_dir, "repo-clone")
        if not os.path.exists(clone_path):
            raise RuntimeError(f"Clone path not found: {clone_path}")

        workspace_id = uuid.uuid4()
        parent_path = os.path.join(self.BASE_PATH, str(workspace_id))
        repo_dir_name = self._repo_directory_name(repo)
        target_path = os.path.join(parent_path, repo_dir_name)

        os.makedirs(parent_path, exist_ok=True)
        with contextlib.suppress(OSError):
            os.chmod(parent_path, 0o777)

        if os.path.exists(target_path):
            shutil.rmtree(target_path, ignore_errors=True)

        # Move the existing clone into the repo-named subdirectory so the
        # workspace path matches the natural git clone layout.
        shutil.move(clone_path, target_path)
        self._make_world_writable(target_path)

        workspace = Workspace(
            id=workspace_id,
            name=name,
            repo_id=repo.id,
            user_id=instance.owner_id,
            branch=instance.branch or "main",
            path=target_path,
            status="ready",
            last_sync_at=datetime.now(),
        )
        session.add(workspace)
        await session.flush()
        return workspace