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

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

from muse.core.paths import muse_dir as _muse_dir, heads_dir as _heads_dir, ref_path as _ref_path, worktrees_dir
from muse.core.types import load_json_file
from muse.core.object_store import restore_object
from muse.core.io import write_text_atomic
from muse.core.refs import (
    get_head_commit_id,
    read_current_branch,
)
from muse.core.commits import read_commit
from muse.core.snapshots import read_snapshot
from muse.core.validation import contain_path, validate_branch_name

logger = logging.getLogger(__name__)


# 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 _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 = load_json_file(meta_path)
        if raw is None:
            logger.warning("⚠️ Could not read worktree metadata for %r", name)
            return None
        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 = _muse_dir(wt_dir)
    # 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 = _muse_dir(repo_root).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 = _muse_dir(repo_root).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 = _ref_path(repo_root, 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 = _muse_dir(wt_path)
    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", _muse_dir(wt_path))
    return repaired