"""Workspace management — compose multiple Muse repositories.

A *workspace* is a collection of related Muse repositories that are developed
together.  Think of a film score that references a sound library, a machine
learning pipeline that includes a dataset repo, or a multi-service codebase
where each service lives in its own Muse repo.

Design
------
Workspaces are distinct from worktrees:

- A **worktree** is one checkout of *one* repo with *one* ``.muse/`` store.
- A **workspace** is an envelope that *links* multiple separate repos together.

The workspace manifest lives at ``.muse/workspace.toml``::

    [[members]]
    name   = "core"
    url    = "https://musehub.ai/acme/core"
    path   = "repos/core"       # relative to workspace root
    branch = "main"             # pinned branch

    [[members]]
    name   = "dataset"
    url    = "https://musehub.ai/acme/dataset"
    path   = "repos/dataset"
    branch = "v2"

Agent workflow
--------------
Each member repo is a fully independent Muse repository.  Agents can commit
to member repos independently and the workspace provides a unified status view
and one-shot sync.

``muse workspace sync`` walks all members and runs ``muse fetch`` + ``muse pull``
so the workspace root always has the latest HEAD for every pinned branch.
``muse workspace sync --workers 8`` parallelises across members.

Security model
--------------
- Manifest size is capped at ``_MAX_MANIFEST_BYTES`` before reading.
- The manifest file and its parent directory are checked for symlinks before
  any read or write to prevent path-traversal attacks.
- All free-form string fields (name, url, path, branch) are TOML-escaped
  before serialisation to prevent injection via crafted values.
- Member ``path`` values are validated to resolve *within* the workspace root.
- Member ``url`` values are checked for a valid scheme (https, http, or local
  path); shell metacharacters are rejected before passing to subprocess.
"""

import concurrent.futures
import logging
import pathlib
import subprocess
from dataclasses import dataclass
from typing import TypedDict

from muse.core.paths import muse_dir as _muse_dir, workspace_toml_path as _workspace_toml_path, shelf_json_path as _shelf_json_path
from muse.core.types import load_json_file
from muse.core.refs import iter_branch_refs

logger = logging.getLogger(__name__)


# 1 MiB — a manifest with 1 000 members at ~200 bytes each is ~200 KiB.
_MAX_MANIFEST_BYTES = 1 * 1024 * 1024

# Allowed URL schemes for member repositories.
_ALLOWED_SCHEMES = frozenset({"https", "http"})

# ---------------------------------------------------------------------------
# Types
# ---------------------------------------------------------------------------

class WorkspaceMemberDict(TypedDict):
    """One entry in the workspace manifest."""

    name: str
    url: str
    path: str
    branch: str

class WorkspaceManifestDict(TypedDict):
    """Top-level workspace manifest."""

    members: list[WorkspaceMemberDict]

@dataclass
class WorkspaceMemberStatus:
    """Runtime status of one workspace member.

    ``branch`` is the configured tracking branch from *workspace.toml* — the
    branch this member is *supposed* to be on according to the manifest.

    ``actual_branch`` is the branch currently checked out in the working
    directory (read from ``HEAD``).  When it differs from ``branch`` the member
    is checked out somewhere unexpected; agents should surface this discrepancy.

    ``head_commit`` is the commit that ``HEAD`` currently resolves to — i.e.
    the actual checked-out commit, not the tip of the configured branch.

    ``shelf_count`` is the number of shelved changesets (0 = nothing on the shelf).

    ``feature_branches`` lists every local branch that is not ``main`` or
    ``dev`` — short-lived task/feat/bugfix branches that have not been cleaned
    up yet.
    """

    name: str
    path: pathlib.Path
    branch: str            # configured tracking branch from workspace.toml
    url: str
    present: bool
    head_commit: str | None   # actual HEAD commit (what HEAD resolves to)
    dirty: bool
    actual_branch: str | None  # currently checked-out branch
    shelf_count: int           # number of shelved changesets
    feature_branches: list[str]  # local branches other than main / dev

class WorkspaceSyncResult(TypedDict):
    """Result of syncing one workspace member.

    ``status`` is one of ``'cloned'``, ``'pulled'``, ``'skipped'``, or
    ``'error: <message>'``.
    """

    name: str
    status: str

# ---------------------------------------------------------------------------
# TOML helpers
# ---------------------------------------------------------------------------

def _toml_escape(value: str) -> str:
    """Escape *value* for safe embedding inside a TOML double-quoted string.

    TOML basic strings forbid unescaped backslash, double-quote, and control
    characters (newline, carriage-return, tab, etc.).  All are escaped here so
    that crafted values like ``core"\\nname = "injected`` or values containing
    literal newlines cannot break the manifest structure.
    """
    return (
        value
        .replace("\\", "\\\\")
        .replace('"', '\\"')
        .replace("\n", "\\n")
        .replace("\r", "\\r")
        .replace("\t", "\\t")
        .replace("\x00", "\\u0000")
    )

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------

def _workspace_path(repo_root: pathlib.Path) -> pathlib.Path:
    return _workspace_toml_path(repo_root)

def find_workspace_root(start: pathlib.Path | None = None) -> pathlib.Path | None:
    """Walk up from *start* (default: cwd) to find the directory containing
    ``.muse/workspace.toml``.  Returns ``None`` if no workspace is found.

    This mirrors ``find_repo_root()`` so that workspace commands resolve the
    correct manifest regardless of CWD or ``-C`` flag usage.
    """
    current = (start or pathlib.Path.cwd()).resolve()
    for directory in (current, *current.parents):
        if _workspace_toml_path(directory).exists():
            return directory
    return None

def require_workspace_root(start: pathlib.Path | None = None) -> pathlib.Path:
    """Return the workspace root or exit with a clear error message."""
    from muse.core.errors import ExitCode
    root = find_workspace_root(start)
    if root is None:
        import sys
        print(
            "❌ Not inside a Muse workspace.\n"
            "   No .muse/workspace.toml found in this directory or any parent.",
            file=sys.stderr,
        )
        raise SystemExit(ExitCode.REPO_NOT_FOUND)
    return root

# ---------------------------------------------------------------------------
# Persistence
# ---------------------------------------------------------------------------

def _load_manifest(repo_root: pathlib.Path) -> WorkspaceManifestDict | None:
    """Read and parse the workspace manifest.

    Security guards applied before any read:

    - Symlink check: a symlink at the manifest path could redirect reads to
      sensitive files outside the repo.
    - Size cap (``_MAX_MANIFEST_BYTES``): a corrupt or tampered manifest cannot
      exhaust memory.
    """
    import tomllib

    path = _workspace_path(repo_root)
    if not path.exists():
        return None
    if path.is_symlink():
        logger.warning(
            "⚠️ Workspace manifest is a symlink — ignoring to prevent path traversal"
        )
        return None
    try:
        size = path.stat().st_size
        if size > _MAX_MANIFEST_BYTES:
            logger.warning(
                "⚠️ Workspace manifest is %.1f MiB — exceeds cap of %d MiB; ignoring",
                size / (1024 * 1024),
                _MAX_MANIFEST_BYTES // (1024 * 1024),
            )
            return None
        raw = tomllib.loads(path.read_text(encoding="utf-8"))
    except Exception as exc:
        logger.warning("⚠️ Could not read workspace manifest: %s", exc)
        return None
    members: list[WorkspaceMemberDict] = []
    for m in raw.get("members", []):
        if not isinstance(m, dict):
            continue
        members.append(
            WorkspaceMemberDict(
                name=str(m.get("name", "")),
                url=str(m.get("url", "")),
                path=str(m.get("path", "")),
                branch=str(m.get("branch", "main")),
            )
        )
    return WorkspaceManifestDict(members=members)

def _save_manifest(repo_root: pathlib.Path, manifest: WorkspaceManifestDict) -> None:
    """Write the manifest atomically.

    Security guards:

    - The manifest file and its parent directory are checked for symlinks
      before writing to prevent path-traversal via a planted symlink.
    - All string values are TOML-escaped to prevent injection.
    """
    path = _workspace_path(repo_root)
    parent = path.parent
    parent.mkdir(parents=True, exist_ok=True)

    if parent.is_symlink():
        raise OSError(f"Refusing to write manifest — parent directory is a symlink: {parent}")
    if path.exists() and path.is_symlink():
        raise OSError(f"Refusing to write manifest — file is a symlink: {path}")

    lines: list[str] = []
    for m in manifest["members"]:
        lines.append("[[members]]")
        lines.append(f'name   = "{_toml_escape(m["name"])}"')
        lines.append(f'url    = "{_toml_escape(m["url"])}"')
        lines.append(f'path   = "{_toml_escape(m["path"])}"')
        lines.append(f'branch = "{_toml_escape(m["branch"])}"')
        lines.append("")
    tmp = path.with_suffix(".tmp")
    tmp.write_text("\n".join(lines), encoding="utf-8")
    tmp.replace(path)

# ---------------------------------------------------------------------------
# Validation helpers
# ---------------------------------------------------------------------------

def _validate_member_name(name: str) -> None:
    """Raise ``ValueError`` if *name* is not a safe workspace member name.

    Allowed: alphanumerics, hyphens, underscores, dots.  No slashes, nulls,
    or shell metacharacters.  Must be 1–64 characters.
    """
    import re
    if not name or len(name) > 64:
        raise ValueError(f"Member name must be 1–64 characters, got {len(name)!r}.")
    if not re.fullmatch(r"[A-Za-z0-9._-]+", name):
        raise ValueError(
            f"Member name {name!r} contains invalid characters. "
            "Use only alphanumerics, hyphens, underscores, and dots."
        )

def _validate_member_url(url: str) -> None:
    """Raise ``ValueError`` if *url* is not a safe member URL or local path.

    Accepted forms:
    - ``https://`` or ``http://`` — remote MuseHub URL.
    - An absolute local path (no scheme).
    - A relative local path (no scheme).

    Rejected:
    - Null bytes in the URL string.
    - ``file://`` — use a bare path instead.
    - Any other scheme (``ftp://``, ``ssh://``, etc.).
    """
    import urllib.parse
    if "\x00" in url:
        raise ValueError("Member URL must not contain null bytes.")
    parsed = urllib.parse.urlparse(url)
    if parsed.scheme and parsed.scheme not in _ALLOWED_SCHEMES:
        raise ValueError(
            f"Member URL scheme {parsed.scheme!r} is not allowed. "
            "Use https://, http://, or a bare filesystem path."
        )

def _validate_member_path(repo_root: pathlib.Path, relative_path: str) -> None:
    """Raise ``ValueError`` if *relative_path* escapes the workspace root.

    Path components like ``../../etc`` would let a crafted manifest point
    members at arbitrary directories.  We resolve the candidate path and
    confirm it sits within *repo_root*.
    """
    if "\x00" in relative_path:
        raise ValueError("Member path must not contain null bytes.")
    candidate = (repo_root / relative_path).resolve()
    try:
        candidate.relative_to(repo_root.resolve())
    except ValueError:
        raise ValueError(
            f"Member path {relative_path!r} resolves outside the workspace root."
        )

# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def add_workspace_member(
    repo_root: pathlib.Path,
    name: str,
    url: str,
    path: str = "",
    branch: str = "main",
) -> None:
    """Register a new member repository in the workspace manifest.

    Args:
        repo_root:  The workspace root (where ``.muse/`` lives).
        name:       Short identifier for this member (alphanumeric, hyphens,
                    underscores, dots; max 64 chars).
        url:        Remote URL (https/http) or local path to the member repo.
        path:       Relative checkout path inside the workspace (default:
                    ``repos/<name>``).  Must not escape the workspace root.
        branch:     Branch to track (default: ``main``).

    Raises:
        ValueError: If name is invalid, URL scheme is disallowed, path escapes
                    the workspace root, or a member with the same name exists.
    """
    from muse.core.validation import validate_branch_name

    _validate_member_name(name)
    _validate_member_url(url)
    validate_branch_name(branch)

    effective_path = path or f"repos/{name}"
    _validate_member_path(repo_root, effective_path)

    manifest = _load_manifest(repo_root) or WorkspaceManifestDict(members=[])
    for m in manifest["members"]:
        if m["name"] == name:
            raise ValueError(f"Workspace member '{name}' already exists.")

    manifest["members"].append(
        WorkspaceMemberDict(
            name=name,
            url=url,
            path=effective_path,
            branch=branch,
        )
    )
    _save_manifest(repo_root, manifest)

def update_workspace_member(
    repo_root: pathlib.Path,
    name: str,
    url: str | None = None,
    path: str | None = None,
    branch: str | None = None,
) -> None:
    """Update the URL, path, or branch for an existing workspace member.

    Only the supplied keyword arguments are changed.  Raises ``ValueError`` if
    no member with *name* exists.

    Args:
        repo_root: The workspace root.
        name:      Member name to update.
        url:       New URL (or ``None`` to keep current).
        path:      New relative checkout path (or ``None`` to keep current).
        branch:    New branch to track (or ``None`` to keep current).

    Raises:
        ValueError: If the member does not exist or any new value is invalid.
    """
    from muse.core.validation import validate_branch_name

    if url is not None:
        _validate_member_url(url)
    if branch is not None:
        validate_branch_name(branch)
    if path is not None:
        _validate_member_path(repo_root, path)

    manifest = _load_manifest(repo_root)
    if manifest is not None:
        for m in manifest["members"]:
            if m["name"] == name:
                if url is not None:
                    m["url"] = url
                if path is not None:
                    m["path"] = path
                if branch is not None:
                    m["branch"] = branch
                _save_manifest(repo_root, manifest)
                return
    raise ValueError(f"Workspace member '{name}' not found.")

def remove_workspace_member(repo_root: pathlib.Path, name: str) -> None:
    """Remove a member from the workspace manifest.

    Does **not** delete the member's directory — only its registration in the
    manifest is removed.

    Raises:
        ValueError: If no member with that name exists.
    """
    manifest = _load_manifest(repo_root)
    if manifest is None:
        raise ValueError("No workspace manifest found.")
    before = len(manifest["members"])
    manifest["members"] = [m for m in manifest["members"] if m["name"] != name]
    if len(manifest["members"]) == before:
        raise ValueError(f"Workspace member '{name}' not found.")
    _save_manifest(repo_root, manifest)

def get_workspace_member(
    repo_root: pathlib.Path,
    name: str,
) -> WorkspaceMemberStatus:
    """Return the status for a single named workspace member.

    Raises:
        ValueError: If no member with that name is registered.
    """
    manifest = _load_manifest(repo_root)
    if manifest is None:
        raise ValueError("No workspace manifest found.")
    for m in manifest["members"]:
        if m["name"] == name:
            return _member_status(repo_root, m)
    raise ValueError(f"Workspace member '{name}' not found.")

def _member_status(repo_root: pathlib.Path, m: WorkspaceMemberDict) -> WorkspaceMemberStatus:
    """Build a ``WorkspaceMemberStatus`` for one manifest entry."""
    import json as _json

    member_path = repo_root / m["path"]
    present = member_path.exists() and (_muse_dir(member_path)).exists()
    head_commit: str | None = None
    dirty = False
    actual_branch: str | None = None
    shelf_count = 0
    feature_branches: list[str] = []

    if present:
        # One subprocess: muse status gives us dirty, actual branch, and HEAD commit.
        try:
            result = subprocess.run(
                ["muse", "status", "--json"],
                capture_output=True,
                text=True,
                cwd=str(member_path),
                timeout=10,
            )
            if result.returncode == 0:
                status_data = _json.loads(result.stdout)
                dirty = bool(status_data.get("dirty", False))
                actual_branch = status_data.get("branch") or None
                head_commit = status_data.get("head_commit") or None
        except Exception as exc:
            logger.debug("Could not read status for member %r: %s", m["name"], exc)

        # Shelf count — read .muse/shelf.json directly; no subprocess needed.
        try:
            shelf_file = _muse_dir(member_path) / "shelf.json"
            if shelf_file.is_file():
                shelf_list = load_json_file(shelf_file)
                if isinstance(shelf_list, list):
                    shelf_count = len(shelf_list)
        except Exception as exc:
            logger.debug("Could not read shelf count for member %r: %s", m["name"], exc)

        # Feature branches — pure file I/O, no subprocess.
        try:
            standard = {"main", "dev"}
            feature_branches = sorted(
                name for name, _ in iter_branch_refs(member_path)
                if name not in standard
            )
        except Exception as exc:
            logger.debug("Could not read branches for member %r: %s", m["name"], exc)

    return WorkspaceMemberStatus(
        name=m["name"],
        path=member_path,
        branch=m["branch"],
        url=m["url"],
        present=present,
        head_commit=head_commit,
        dirty=dirty,
        actual_branch=actual_branch,
        shelf_count=shelf_count,
        feature_branches=feature_branches,
    )

def list_workspace_members(repo_root: pathlib.Path) -> list[WorkspaceMemberStatus]:
    """Return status for every workspace member.

    Each member's status requires one ``muse status`` subprocess call plus
    several file-I/O reads.  Members are processed concurrently via a
    thread pool so the total wall time is bounded by the slowest single
    member rather than the sum across all members.

    Results are returned in the same order as the manifest.
    """
    manifest = _load_manifest(repo_root)
    if manifest is None:
        return []
    members = manifest["members"]
    if len(members) <= 1:
        return [_member_status(repo_root, m) for m in members]

    results: list[WorkspaceMemberStatus | None] = [None] * len(members)
    with concurrent.futures.ThreadPoolExecutor(max_workers=len(members)) as pool:
        future_to_idx = {
            pool.submit(_member_status, repo_root, m): i
            for i, m in enumerate(members)
        }
        for future in concurrent.futures.as_completed(future_to_idx):
            idx = future_to_idx[future]
            try:
                results[idx] = future.result()
            except Exception as exc:
                logger.warning(
                    "⚠️ Could not read status for member %r: %s",
                    members[idx]["name"], exc,
                )
    return [r for r in results if r is not None]

def sync_workspace_member(
    repo_root: pathlib.Path,
    member: WorkspaceMemberDict,
    dry_run: bool = False,
) -> WorkspaceSyncResult:
    """Clone or pull the latest state for one workspace member.

    Returns a :class:`WorkspaceSyncResult` dict with ``name`` and ``status``.
    ``status`` is one of ``'cloned'``, ``'pulled'``, ``'skipped'`` (dry-run),
    or ``'error: <message>'``.

    Security: ``url`` and ``branch`` are passed as separate list elements to
    ``subprocess.run`` (never via the shell), so shell injection is not
    possible.  Size of error output is capped at 200 chars.
    """
    member_path = repo_root / member["path"]
    name = member["name"]

    if dry_run:
        action = "clone" if (not member_path.exists() or not (_muse_dir(member_path)).exists()) else "pull"
        return WorkspaceSyncResult(name=name, status=f"skipped (dry-run would {action})")

    if not member_path.exists() or not (_muse_dir(member_path)).exists():
        member_path.parent.mkdir(parents=True, exist_ok=True)
        result = subprocess.run(
            ["muse", "clone", member["url"], str(member_path)],
            capture_output=True,
            text=True,
            timeout=300,
        )
        if result.returncode != 0:
            err = result.stderr.strip()[:200]
            logger.warning("⚠️ Clone failed for member %r: %s", name, err)
            return WorkspaceSyncResult(name=name, status=f"error: {err}")
        return WorkspaceSyncResult(name=name, status="cloned")

    result = subprocess.run(
        ["muse", "pull", "--branch", member["branch"]],
        capture_output=True,
        text=True,
        cwd=str(member_path),
        timeout=120,
    )
    if result.returncode != 0:
        err = result.stderr.strip()[:200]
        logger.warning("⚠️ Pull failed for member %r: %s", name, err)
        return WorkspaceSyncResult(name=name, status=f"error: {err}")
    return WorkspaceSyncResult(name=name, status="pulled")

def sync_workspace(
    repo_root: pathlib.Path,
    member_name: str | None = None,
    dry_run: bool = False,
    workers: int = 1,
) -> list[WorkspaceSyncResult]:
    """Sync all (or one named) workspace members.

    Args:
        repo_root:   The workspace root.
        member_name: Sync only this member (default: all).
        dry_run:     Show what would happen without doing it.
        workers:     Number of parallel sync workers (default: 1 — sequential).
                     Set to > 1 to parallelise across members.

    Returns:
        List of :class:`WorkspaceSyncResult` dicts, one per member synced.
    """
    manifest = _load_manifest(repo_root)
    if manifest is None:
        return []

    targets = (
        [m for m in manifest["members"] if m["name"] == member_name]
        if member_name is not None
        else manifest["members"]
    )

    if workers <= 1 or len(targets) <= 1:
        return [sync_workspace_member(repo_root, m, dry_run=dry_run) for m in targets]

    results: list[WorkspaceSyncResult] = []
    effective_workers = min(workers, len(targets))
    with concurrent.futures.ThreadPoolExecutor(max_workers=effective_workers) as pool:
        futures = {
            pool.submit(sync_workspace_member, repo_root, m, dry_run): m["name"]
            for m in targets
        }
        for future in concurrent.futures.as_completed(futures):
            try:
                results.append(future.result())
            except Exception as exc:
                name = futures[future]
                logger.warning("⚠️ Unexpected sync error for member %r: %s", name, exc)
                results.append(WorkspaceSyncResult(name=name, status=f"error: {exc}"))
    return results