"""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.
"""

from __future__ import annotations

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

logger = logging.getLogger(__name__)

_WORKSPACE_FILE = ".muse/workspace.toml"

# 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.

    ``stash_count`` is the number of stashed changesets (0 = clean).

    ``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
    stash_count: int           # number of stashed 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 repo_root / ".muse" / "workspace.toml"


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 (directory / ".muse" / "workspace.toml").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 (member_path / ".muse").exists()
    head_commit: str | None = None
    dirty = False
    actual_branch: str | None = None
    stash_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)

        # Stash count — one subprocess.
        try:
            result = subprocess.run(
                ["muse", "stash", "list", "--json"],
                capture_output=True,
                text=True,
                cwd=str(member_path),
                timeout=10,
            )
            if result.returncode == 0:
                stash_list = _json.loads(result.stdout)
                if isinstance(stash_list, list):
                    stash_count = len(stash_list)
        except Exception as exc:
            logger.debug("Could not read stash list for member %r: %s", m["name"], exc)

        # Feature branches — pure file I/O, no subprocess.
        try:
            heads_dir = member_path / ".muse" / "refs" / "heads"
            if heads_dir.is_dir():
                standard = {"main", "dev"}
                feature_branches = sorted(
                    p.name
                    for p in heads_dir.iterdir()
                    if p.is_file() and p.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,
        stash_count=stash_count,
        feature_branches=feature_branches,
    )


def list_workspace_members(repo_root: pathlib.Path) -> list[WorkspaceMemberStatus]:
    """Return status for every workspace member."""
    manifest = _load_manifest(repo_root)
    if manifest is None:
        return []
    return [_member_status(repo_root, m) for m in manifest["members"]]


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 (member_path / ".muse").exists()) else "pull"
        return WorkspaceSyncResult(name=name, status=f"skipped (dry-run would {action})")

    if not member_path.exists() or not (member_path / ".muse").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