"""Workspace Manager - Handles git worktrees and Docker containers for isolated execution."""

import os
import shutil
from pathlib import Path

import docker
import git


class WorkspaceError(Exception):
    """Base exception for workspace operations."""


class GitWorktreeError(WorkspaceError):
    """Exception for git worktree failures."""


class DockerProvisionError(WorkspaceError):
    """Exception for Docker provisioning failures."""


class WorkspaceManager:
    """Manages git worktrees and Docker containers for isolated Dev Agent execution."""

    def __init__(self, repo_path: str, docker_image: str = "python:3.11-slim"):
        """Initialize WorkspaceManager.

        Args:
            repo_path: Path to the git repository.
            docker_image: Docker image to use for clean room containers.
        """
        try:
            self.repo = git.Repo(repo_path)
        except git.InvalidGitRepositoryError as e:
            raise GitWorktreeError(f"Invalid git repository: {repo_path}") from e
        except git.NoSuchPathError as e:
            raise GitWorktreeError(f"Repository path not found: {repo_path}") from e

        self.repo_path = Path(repo_path).resolve()
        self.docker_image = docker_image
        self.active_workspaces: dict[str, dict] = {}

        try:
            self.docker_client = docker.from_env()
        except docker.errors.DockerException as e:
            raise DockerProvisionError(
                "Failed to connect to Docker daemon. Is Docker running?"
            ) from e

    async def create_worktree(self, task_id: str, base_branch: str = "main") -> str:
        """Create a git worktree for a task.

        Args:
            task_id: Unique identifier for the task.
            base_branch: Branch to base the worktree on.

        Returns:
            Absolute path to the created worktree.

        Raises:
            GitWorktreeError: If worktree creation fails.
        """
        branch_name = f"feature/task-{task_id}"
        worktree_path = str(self.repo_path.parent / "worktrees" / task_id)

        # Validate base branch exists
        try:
            self.repo.git.rev_parse("--verify", base_branch)
        except git.GitCommandError as e:
            raise GitWorktreeError(
                f"Base branch '{base_branch}' does not exist"
            ) from e

        # Check if worktree path already exists
        if os.path.exists(worktree_path):
            raise GitWorktreeError(
                f"Worktree path already exists: {worktree_path}"
            )

        # Check if branch already exists
        if branch_name in [ref.name for ref in self.repo.branches]:
            raise GitWorktreeError(
                f"Branch already exists: {branch_name}"
            )

        try:
            os.makedirs(os.path.dirname(worktree_path), exist_ok=True)
            self.repo.git.worktree(
                "add", worktree_path, "-b", branch_name, base_branch
            )
        except git.GitCommandError as e:
            raise GitWorktreeError(
                f"Failed to create worktree for task {task_id}: {e}"
            ) from e

        return str(Path(worktree_path).resolve())

    async def spin_up_clean_room(self, worktree_path: str, task_id: str):
        """Create an isolated Docker container for a task.

        Args:
            worktree_path: Path to the git worktree to mount.
            task_id: Unique identifier for the task.

        Returns:
            Container object with metadata.

        Raises:
            DockerProvisionError: If container creation fails.
        """
        try:
            self.docker_client.images.pull(self.docker_image)
        except docker.errors.APIError as e:
            raise DockerProvisionError(
                f"Failed to pull image '{self.docker_image}': {e}"
            ) from e

        try:
            container = self.docker_client.containers.create(
                image=self.docker_image,
                name=f"appfactory-task-{task_id}",
                volumes={
                    worktree_path: {"bind": "/workspace", "mode": "rw"}
                },
                working_dir="/workspace",
                network_mode="none",
                auto_remove=False,
                detach=True,
                command="sleep infinity",
            )
        except docker.errors.APIError as e:
            raise DockerProvisionError(
                f"Failed to create container for task {task_id}: {e}"
            ) from e

        self.active_workspaces[task_id] = {
            "task_id": task_id,
            "worktree_path": worktree_path,
            "container_id": container.id,
            "container": container,
        }

        return container

    async def cleanup_workspace(self, task_id: str, container=None):
        """Clean up a workspace by removing its container and worktree.

        Args:
            task_id: Unique identifier for the task.
            container: Optional container object. If None, uses the registered one.

        Raises:
            WorkspaceError: If cleanup fails completely.
        """
        workspace = self.active_workspaces.get(task_id, {})
        errors = []

        # Resolve container
        if container is None:
            container = workspace.get("container")

        # Stop and remove container
        if container is not None:
            try:
                container.stop(timeout=5)
            except Exception:
                pass  # Container may already be stopped
            try:
                container.remove(force=True)
            except Exception as e:
                errors.append(f"Container removal failed: {e}")

        # Remove worktree
        worktree_path = workspace.get("worktree_path")
        if worktree_path is None:
            worktree_path = str(self.repo_path.parent / "worktrees" / task_id)

        try:
            self.repo.git.worktree("remove", worktree_path, "--force")
        except git.GitCommandError:
            # Worktree may already be removed; try cleaning up the directory
            if os.path.exists(worktree_path):
                try:
                    shutil.rmtree(worktree_path)
                except OSError as e:
                    errors.append(f"Worktree directory removal failed: {e}")

        # Prune worktree references
        try:
            self.repo.git.worktree("prune")
        except git.GitCommandError:
            pass

        # Remove from registry
        self.active_workspaces.pop(task_id, None)

        if errors:
            raise WorkspaceError(
                f"Cleanup completed with errors for task {task_id}: {'; '.join(errors)}"
            )

    def get_active_workspaces(self) -> list:
        """Return list of active workspace info dicts.

        Returns:
            List of dicts with task_id, worktree_path, and container_id.
        """
        return [
            {
                "task_id": info["task_id"],
                "worktree_path": info["worktree_path"],
                "container_id": info["container_id"],
            }
            for info in self.active_workspaces.values()
        ]

    async def cleanup_all(self):
        """Cleanup all active workspaces. Used for graceful shutdown."""
        task_ids = list(self.active_workspaces.keys())
        errors = []
        for task_id in task_ids:
            try:
                await self.cleanup_workspace(task_id)
            except WorkspaceError as e:
                errors.append(str(e))
        if errors:
            raise WorkspaceError(
                f"Cleanup all completed with errors: {'; '.join(errors)}"
            )
