"""Worktree management — multiple simultaneous branch checkouts.

A *worktree* is a second (or third, …) checked-out working tree linked to
the same ``.muse/`` repository.  Each worktree has its own branch and its own
working directory, so multiple agents — or multiple human engineers — can
work on different branches simultaneously without interfering with each other.

Layout
------
Each linked worktree lives in a sibling directory of the repository root::

    myproject/              ← main worktree  (holds .muse/)
      .muse/
        worktrees/
          <name>.json       ← metadata for each linked worktree
          <name>.HEAD       ← HEAD ref for the linked worktree

    myproject-<name>/       ← linked worktree directory
                            ← worktree working directory

The shared ``.muse/`` directory is the single source of truth for commits,
snapshots, objects, and branch refs.  Each worktree has its own HEAD file
stored inside the main ``.muse/worktrees/<name>.HEAD``.

Security model
--------------
- **Name validation**: Worktree names pass through ``validate_branch_name``
  (no path separators, no null bytes, no control characters).  This ensures
  the derived meta path and HEAD path cannot escape ``.muse/worktrees/``.
- **Symlink guard on meta files**: ``_load_meta`` rejects a symlink at the
  meta file path before any read.
- **Size cap on meta files**: ``_load_meta`` refuses files larger than
  ``_MAX_META_BYTES`` to guard against memory exhaustion.
- **Path safety on delete**: ``remove_worktree`` and ``prune_worktrees``
  call ``_safe_delete_path`` which refuses to ``rmtree`` a path that is a
  symlink itself or that resolves inside the shared ``.muse/`` store.

Agent concurrency
-----------------
Multiple agents can operate on separate worktrees simultaneously.  Each
worktree's HEAD is independent; commits from one worktree appear immediately
in all others (they share the object store).
"""

from __future__ import annotations

import json
import logging
import pathlib
import shutil
from dataclasses import dataclass
from typing import TypedDict

from muse.core.object_store import restore_object
from muse.core.store import (
    get_head_commit_id,
    read_commit,
    read_current_branch,
    read_snapshot,
    write_text_atomic,
)
from muse.core.validation import contain_path, validate_branch_name

logger = logging.getLogger(__name__)

# Directory inside .muse/ that holds per-worktree metadata and HEAD files.
_WORKTREES_META_DIR = ".muse/worktrees"

# Guard against tampered or pathologically large meta files.
_MAX_META_BYTES: int = 4 * 1024  # 4 KiB — more than enough for any meta record


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


class WorktreeRecord(TypedDict):
    """Persisted metadata for a linked worktree."""

    name: str
    branch: str
    path: str  # absolute path to the worktree directory


@dataclass
class WorktreeInfo:
    """Runtime information about a worktree."""

    name: str
    branch: str
    path: pathlib.Path
    head_commit: str | None
    is_main: bool = False


class WorktreeStatusResult(TypedDict):
    """Machine-readable status of a single worktree."""

    name: str
    branch: str
    path: str
    head_commit: str | None
    present: bool
    is_main: bool


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


def _worktrees_dir(repo_root: pathlib.Path) -> pathlib.Path:
    return repo_root / _WORKTREES_META_DIR


def _worktree_meta_path(repo_root: pathlib.Path, name: str) -> pathlib.Path:
    return _worktrees_dir(repo_root) / f"{name}.json"


def _worktree_head_path(repo_root: pathlib.Path, name: str) -> pathlib.Path:
    return _worktrees_dir(repo_root) / f"{name}.HEAD"


def _worktree_dir(repo_root: pathlib.Path, name: str) -> pathlib.Path:
    """Return the default path of the linked worktree directory (sibling of repo_root)."""
    parent = repo_root.parent
    repo_name = repo_root.name
    return parent / f"{repo_name}-{name}"


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _load_meta(repo_root: pathlib.Path, name: str) -> WorktreeRecord | None:
    """Read and validate the metadata file for *name*.

    Safety guards applied before any read:

    - **Symlink check**: a symlink at the meta path could redirect writes to
      arbitrary locations or reads to sensitive files.
    - **Size cap** (``_MAX_META_BYTES``): a tampered or corrupt meta file
      cannot be used to exhaust memory.
    """
    meta_path = _worktree_meta_path(repo_root, name)
    if not meta_path.exists():
        return None
    if meta_path.is_symlink():
        logger.warning("⚠️ Worktree meta file for %r is a symlink — ignoring", name)
        return None
    try:
        if meta_path.stat().st_size > _MAX_META_BYTES:
            logger.warning(
                "⚠️ Worktree meta file for %r exceeds size cap (%d bytes) — ignoring",
                name,
                _MAX_META_BYTES,
            )
            return None
        raw = json.loads(meta_path.read_text(encoding="utf-8"))
        return WorktreeRecord(
            name=str(raw["name"]),
            branch=str(raw["branch"]),
            path=str(raw["path"]),
        )
    except (KeyError, ValueError, OSError) as exc:
        logger.warning("⚠️ Could not read worktree metadata for %r: %s", name, exc)
        return None


def _save_meta(repo_root: pathlib.Path, record: WorktreeRecord) -> None:
    """Write *record* to the metadata file atomically."""
    meta_path = _worktree_meta_path(repo_root, record["name"])
    write_text_atomic(meta_path, json.dumps(record, indent=2))


def _write_worktree_pointer(wt_dir: pathlib.Path, repo_root: pathlib.Path) -> None:
    """Write a ``.muse`` pointer file in *wt_dir* pointing to *repo_root*'s store.

    The file contains a single line::

        musestore: /absolute/path/to/main/.muse

    This mirrors git's ``.git`` worktree file, enabling ``find_repo_root``
    to resolve the shared object store from any worktree directory.

    No-ops when ``wt_dir`` is the main repo itself or when ``.muse`` already
    exists as a directory (a legacy worktree with its own store).
    """
    pointer_path = wt_dir / ".muse"
    # Never clobber a real .muse/ store directory.
    if pointer_path.is_dir():
        logger.debug("Skipping pointer write — %s is already a store directory", pointer_path)
        return
    store_path = (repo_root / ".muse").resolve()
    # Don't write a self-referential pointer (wt_dir IS the main repo).
    if pointer_path.resolve() == store_path:
        logger.debug("Skipping pointer write — wt_dir is the main repo root")
        return
    write_text_atomic(pointer_path, f"musestore: {store_path}\n")
    logger.debug("Wrote worktree pointer %s → %s", pointer_path, store_path)


def _safe_delete_path(repo_root: pathlib.Path, path: pathlib.Path) -> bool:
    """Delete *path* and its contents, with safety guards.

    Refuses deletion when:

    - *path* is a symlink (could target an unrelated directory).
    - *path* resolves to be inside the shared ``.muse/`` store — deleting it
      would corrupt the repository.
    - *path* does not exist (no-op, returns True so callers can proceed).

    Returns:
        ``True`` if the directory was deleted or did not exist.
        ``False`` if the deletion was refused for safety reasons.
    """
    if not path.exists():
        return True
    if path.is_symlink():
        logger.warning(
            "⚠️ Refusing to delete worktree path %s — it is a symlink", path
        )
        return False
    muse_dir = (repo_root / ".muse").resolve()
    try:
        resolved = path.resolve()
    except OSError:
        logger.warning("⚠️ Could not resolve worktree path %s", path)
        return False
    try:
        resolved.relative_to(muse_dir)
        # Path is inside .muse/ — refuse.
        logger.warning(
            "⚠️ Refusing to delete worktree path %s — it resolves inside .muse/", path
        )
        return False
    except ValueError:
        pass  # Not inside .muse/ — safe to delete.
    shutil.rmtree(path)
    return True


def _read_main_branch(repo_root: pathlib.Path) -> str:
    return read_current_branch(repo_root)


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


def add_worktree(
    repo_root: pathlib.Path,
    name: str,
    branch: str,
    path: pathlib.Path | None = None,
) -> pathlib.Path:
    """Create and populate a new linked worktree.

    Args:
        repo_root:  Main repository root (where ``.muse/`` lives).
        name:       Short identifier for the worktree (validated like a branch
                    name — no path separators, no null bytes).
        branch:     Branch to check out in the new worktree.
        path:       Explicit filesystem path for the worktree directory.  When
                    ``None`` (default) the standard sibling layout is used:
                    ``<repo_root.parent>/<repo_root.name>-<name>``.

    Returns:
        The path to the newly created worktree directory.

    Raises:
        ValueError: If the name is invalid, the worktree already exists, or
                    the branch does not exist.
    """
    validate_branch_name(name)

    wt_dir = path if path is not None else _worktree_dir(repo_root, name)
    meta_path = _worktree_meta_path(repo_root, name)

    if meta_path.exists():
        raise ValueError(f"Worktree '{name}' already exists.")
    if wt_dir.exists():
        raise ValueError(f"Directory '{wt_dir}' already exists.")

    # Verify the branch exists.
    branch_ref = repo_root / ".muse" / "refs" / "heads" / branch
    if not branch_ref.exists():
        raise ValueError(f"Branch '{branch}' does not exist.")

    # Create the worktree directory.
    wt_dir.mkdir(parents=True)

    # Write .muse pointer file so `muse` commands work inside the worktree.
    _write_worktree_pointer(wt_dir, repo_root)

    # Write the worktree HEAD file.
    head_path = _worktree_head_path(repo_root, name)
    write_text_atomic(head_path, f"refs/heads/{branch}\n")

    # Populate the worktree from the branch's latest snapshot.
    commit_id = get_head_commit_id(repo_root, branch)
    if commit_id:
        commit = read_commit(repo_root, commit_id)
        if commit:
            snap = read_snapshot(repo_root, commit.snapshot_id)
            if snap:
                for rel_path, object_id in snap.manifest.items():
                    try:
                        dest = contain_path(wt_dir, rel_path)
                    except ValueError as exc:
                        logger.warning("⚠️ Skipping unsafe path %r: %s", rel_path, exc)
                        continue
                    restore_object(repo_root, object_id, dest)

    # Persist metadata.
    record: WorktreeRecord = {
        "name": name,
        "branch": branch,
        "path": str(wt_dir),
    }
    _save_meta(repo_root, record)

    logger.info("✅ Worktree '%s' created at %s (branch: %s)", name, wt_dir, branch)
    return wt_dir


def list_worktrees(repo_root: pathlib.Path) -> list[WorktreeInfo]:
    """Return all worktrees (main + linked), sorted by name.

    The main worktree is always first; linked worktrees follow in
    lexicographic order of name.
    """
    results: list[WorktreeInfo] = []

    # Main worktree.
    main_branch = _read_main_branch(repo_root)
    main_head = get_head_commit_id(repo_root, main_branch)
    results.append(WorktreeInfo(
        name="(main)",
        branch=main_branch,
        path=repo_root,
        head_commit=main_head,
        is_main=True,
    ))

    wt_dir = _worktrees_dir(repo_root)
    if not wt_dir.exists():
        return results

    for meta_file in sorted(wt_dir.glob("*.json")):
        name = meta_file.stem
        record = _load_meta(repo_root, name)
        if record is None:
            continue
        wt_path = pathlib.Path(record["path"])
        branch = record["branch"]
        head_path = _worktree_head_path(repo_root, name)
        commit_id = get_head_commit_id(repo_root, branch) if head_path.exists() else None
        results.append(WorktreeInfo(
            name=name,
            branch=branch,
            path=wt_path,
            head_commit=commit_id,
        ))
    return results


def remove_worktree(repo_root: pathlib.Path, name: str, force: bool = False) -> None:
    """Remove a linked worktree.

    The branch itself is not deleted — only the worktree directory and its
    metadata are removed.  Commits already made in the worktree remain in the
    shared object store.

    Args:
        repo_root:  Main repository root.
        name:       Name of the worktree to remove.
        force:      Accepted for interface compatibility.  Currently has no
                    effect since Muse does not track working-tree dirtiness
                    per-worktree.

    Raises:
        ValueError: If the worktree does not exist or its metadata is corrupt.
    """
    validate_branch_name(name)

    meta_path = _worktree_meta_path(repo_root, name)
    if not meta_path.exists():
        raise ValueError(f"Worktree '{name}' does not exist.")

    record = _load_meta(repo_root, name)
    if record is None:
        raise ValueError(f"Could not read metadata for worktree '{name}'.")

    wt_path = pathlib.Path(record["path"])
    if not _safe_delete_path(repo_root, wt_path):
        raise ValueError(
            f"Refusing to delete worktree path '{wt_path}' — "
            "it is a symlink or resolves inside .muse/."
        )

    meta_path.unlink(missing_ok=True)
    head_path = _worktree_head_path(repo_root, name)
    head_path.unlink(missing_ok=True)
    # Belt-and-suspenders: remove the pointer file if it still exists
    # (the directory deletion above should have removed it already).
    pointer_path = wt_path / ".muse"
    pointer_path.unlink(missing_ok=True)

    logger.info("Worktree '%s' removed.", name)


def prune_worktrees(repo_root: pathlib.Path, *, dry_run: bool = False) -> list[str]:
    """Remove metadata for worktrees whose directories no longer exist.

    Args:
        repo_root:  Main repository root.
        dry_run:    When ``True``, report what would be pruned without
                    making any filesystem changes.

    Returns:
        Names of pruned (or would-be-pruned, when *dry_run* is ``True``)
        worktrees.
    """
    pruned: list[str] = []
    wt_dir = _worktrees_dir(repo_root)
    if not wt_dir.exists():
        return pruned
    for meta_file in list(wt_dir.glob("*.json")):
        name = meta_file.stem
        record = _load_meta(repo_root, name)
        if record is None:
            if not dry_run:
                meta_file.unlink(missing_ok=True)
            pruned.append(name)
            continue
        wt_path = pathlib.Path(record["path"])
        if not wt_path.exists():
            if not dry_run:
                meta_file.unlink(missing_ok=True)
                _worktree_head_path(repo_root, name).unlink(missing_ok=True)
            pruned.append(name)
    return pruned


def get_worktree_status(repo_root: pathlib.Path, name: str) -> WorktreeStatusResult:
    """Return the status of a single named worktree.

    Covers both linked worktrees (by name) and the implicit main worktree
    when *name* is ``"(main)"`` or ``"main"``.

    Returns:
        A ``WorktreeStatusResult`` with present/absent flag, current branch,
        and HEAD commit — ready for JSON serialisation by the CLI.

    Raises:
        ValueError: If no worktree with *name* exists.
    """
    # Main worktree shortcut.
    if name in {"(main)", "main"}:
        branch = _read_main_branch(repo_root)
        head = get_head_commit_id(repo_root, branch)
        return WorktreeStatusResult(
            name="(main)",
            branch=branch,
            path=str(repo_root),
            head_commit=head,
            present=repo_root.exists(),
            is_main=True,
        )

    validate_branch_name(name)
    meta_path = _worktree_meta_path(repo_root, name)
    if not meta_path.exists():
        raise ValueError(f"Worktree '{name}' does not exist.")

    record = _load_meta(repo_root, name)
    if record is None:
        raise ValueError(f"Could not read metadata for worktree '{name}'.")

    wt_path = pathlib.Path(record["path"])
    branch = record["branch"]
    head_path = _worktree_head_path(repo_root, name)
    head = get_head_commit_id(repo_root, branch) if head_path.exists() else None
    return WorktreeStatusResult(
        name=name,
        branch=branch,
        path=str(wt_path),
        head_commit=head,
        present=wt_path.exists() and not wt_path.is_symlink(),
        is_main=False,
    )


def repair_worktree_pointers(repo_root: pathlib.Path) -> list[str]:
    """Write missing ``.muse`` pointer files for all registered worktrees.

    Idempotent — safe to run multiple times.  Returns a list of worktree
    names that were repaired (pointer file written or overwritten).
    """
    repaired: list[str] = []
    for wt in list_worktrees(repo_root):
        if wt.is_main:
            continue  # Main worktree has a real .muse/ store — never needs a pointer.
        wt_path = wt.path
        if not wt_path.is_dir():
            continue
        _write_worktree_pointer(wt_path, repo_root)
        repaired.append(wt.name)
        logger.info("✅ Repaired worktree pointer: %s", wt_path / ".muse")
    return repaired