"""``muse pull`` — fetch from a remote and merge into the current branch.

Combines ``muse fetch`` and ``muse merge`` in a single command:

1. Downloads commits, snapshots, and objects from the remote.
2. Updates the remote-tracking pointer.
3. Performs a three-way merge of the remote branch HEAD into the current branch.

If the remote branch is already an ancestor of the local HEAD (fast-forward),
the local branch ref and working tree are advanced without a merge commit.

Pass ``--no-merge`` to stop after the fetch step (equivalent to ``muse fetch``).
Pass ``--ff-only`` to refuse the pull if a fast-forward is not possible.
Pass ``--dry-run`` / ``-n`` to preview what would change without touching the
working tree or writing any commits.

JSON output (``--format json`` / ``--json``) schema::

    {
      "status":           "up_to_date | fast_forward | merged | conflict | fetched | dry_run",
      "remote":           "<name>",
      "branch":           "<remote_branch>",
      "local_branch":     "<local_branch>",
      "commits_received": <N>,
      "objects_written":  <N>,
      "head":             "<sha256> | null",
      "conflict_paths":   ["<path>", ...],
      "dry_run":          false
    }

Exit codes::

    0 — success (up_to_date, fast_forward, merged, fetched, or dry_run)
    1 — remote not configured, branch not found, fetch failed, ff-only refused,
        authentication failure, format error
    2 — merge conflict (requires manual resolution)
"""

from __future__ import annotations

import argparse
import datetime
import json
import logging
import pathlib
import sys
from typing import TypedDict

from muse.cli.config import get_signing_identity, get_remote, get_remote_head, get_upstream, set_remote_head
from muse.core.errors import ExitCode
from muse.core.merge_engine import find_merge_base, write_merge_state
from muse.core.object_store import has_object, write_object
from muse.core._types import Manifest
from muse.core.pack import apply_pack
from muse.core.repo import read_repo_id, require_repo
from muse.core.snapshot import compute_commit_id, compute_snapshot_id, directories_from_manifest
from muse.core.store import (
    CommitRecord,
    SnapshotRecord,
    get_all_commits,
    get_head_commit_id,
    get_head_snapshot_manifest,
    read_commit,
    read_current_branch,
    read_snapshot,
    write_branch_ref,
    write_commit,
    write_snapshot,
)
from muse.core.transport import (
    HttpTransport,
    LocalFileTransport,
    TransportError,
    make_transport,
    negotiate_have,
)
from muse.core.workdir import apply_manifest
from muse.domain import SnapshotManifest, StructuredMergePlugin
from muse.plugins.registry import read_domain, resolve_plugin
from muse.core.validation import sanitize_display
from muse.core._types import Manifest

logger = logging.getLogger(__name__)


class _PullJson(TypedDict):
    """Stable JSON schema emitted by ``muse pull --json``."""

    status: str            # up_to_date | fast_forward | merged | conflict | fetched | dry_run
    remote: str
    branch: str            # remote branch pulled from
    local_branch: str
    commits_received: int
    objects_written: int
    head: str | None       # new local HEAD after the pull, null on conflict/dry_run
    conflict_paths: list[str]
    dry_run: bool


def register(subparsers: "argparse._SubParsersAction[argparse.ArgumentParser]") -> None:
    """Register the ``muse pull`` subcommand and all its flags."""
    parser = subparsers.add_parser(
        "pull",
        help="Fetch from a remote and merge into the current branch.",
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "remote", nargs="?", default="origin",
        help="Remote name to pull from (default: origin).",
    )
    parser.add_argument(
        "branch_pos", nargs="?", default=None, metavar="BRANCH",
        help="Remote branch to pull (default: tracked branch or current branch). Same as --branch.",
    )
    parser.add_argument(
        "--branch", "-b", default=None, dest="branch_flag",
        help="Remote branch to pull (default: tracked branch or current branch).",
    )
    parser.add_argument(
        "--no-merge", action="store_true", dest="no_merge",
        help="Only fetch; do not merge into the current branch.",
    )
    parser.add_argument(
        "--ff-only", action="store_true", dest="ff_only",
        help="Refuse to pull if a fast-forward merge is not possible.",
    )
    parser.add_argument(
        "-n", "--dry-run", action="store_true",
        help="Preview what would change without modifying the working tree or writing commits.",
    )
    parser.add_argument(
        "-m", "--message", default=None,
        help="Override the merge commit message.",
    )
    parser.add_argument(
        "--format", "-f", default="text", dest="fmt",
        help="Output format: text or json.",
    )
    parser.add_argument(
        "--json", action="store_const", const="json", dest="fmt",
        help="Shorthand for --format json.",
    )
    parser.set_defaults(func=run)


def run(args: argparse.Namespace) -> None:
    """Fetch from a remote and merge into the current branch.

    All progress and error messages go to **stderr**.  ``--format json``
    (or ``--json``) emits a single JSON object on stdout for agent pipelines.

    ``--dry-run`` contacts the remote to discover the current HEAD but does
    not fetch objects, apply manifests, or write any commits.  The JSON output
    shows exactly what *would* happen (``status``, ``commits_received``
    estimate, ``conflict_paths`` if a three-way merge would be needed).

    ``--ff-only`` exits with code 1 if the remote HEAD is not a descendant of
    the local HEAD, preventing silent three-way merge commits.

    JSON schema::

        {
          "status":           "up_to_date | fast_forward | merged | conflict | fetched | dry_run",
          "remote":           "<remote_name>",
          "branch":           "<remote_branch>",
          "local_branch":     "<local_branch>",
          "commits_received": <N>,
          "objects_written":  <N>,
          "head":             "<sha256> | null",
          "conflict_paths":   ["<path>", ...],
          "dry_run":          false
        }

    Exit codes::

        0 — success
        1 — configuration, network, or ff-only refusal error
        2 — merge conflict (resolve, then ``muse commit``)
    """
    remote: str = args.remote
    branch: str | None = (
        getattr(args, "branch_flag", None) or getattr(args, "branch_pos", None)
    )
    no_merge: bool = args.no_merge
    ff_only: bool = getattr(args, "ff_only", False)
    dry_run: bool = getattr(args, "dry_run", False)
    message: str | None = args.message
    fmt: str = getattr(args, "fmt", "text")

    if fmt not in ("text", "json"):
        print(
            f"❌ Unknown --format '{sanitize_display(fmt)}'. Choose text or json.",
            file=sys.stderr,
        )
        raise SystemExit(ExitCode.USER_ERROR)

    root = require_repo()

    url = get_remote(remote, root)
    if url is None:
        print(
            f"❌ Remote '{sanitize_display(remote)}' is not configured.\n"
            f"  Add it with: muse remote add {sanitize_display(remote)} <url>",
            file=sys.stderr,
        )
        raise SystemExit(ExitCode.USER_ERROR)

    signing = get_signing_identity(root, remote_url=url)
    current_branch = read_current_branch(root)
    target_branch = branch or get_upstream(current_branch, root) or current_branch

    transport: HttpTransport | LocalFileTransport = make_transport(url)

    # ── Phase 0: discover remote HEAD ────────────────────────────────────────
    try:
        info = transport.fetch_remote_info(url, signing)
    except TransportError as exc:
        print(
            f"❌ Cannot reach remote '{sanitize_display(remote)}': "
            f"{sanitize_display(str(exc))}",
            file=sys.stderr,
        )
        raise SystemExit(ExitCode.INTERNAL_ERROR)

    remote_commit_id = info["branch_heads"].get(target_branch)
    if remote_commit_id is None:
        print(
            f"❌ Branch '{sanitize_display(target_branch)}' does not exist "
            f"on remote '{sanitize_display(remote)}'.",
            file=sys.stderr,
        )
        raise SystemExit(ExitCode.USER_ERROR)

    # ── Phase 1: fetch ────────────────────────────────────────────────────────
    commits_received: int = 0
    objects_written: int = 0

    already_known = get_remote_head(remote, target_branch, root)
    if already_known == remote_commit_id:
        # We have everything for this commit — skip network fetch entirely.
        if fmt == "text":
            print(
                f"✅ Fetched 0 commit(s), 0 new object(s) from "
                f"{sanitize_display(remote)}/{sanitize_display(target_branch)} "
                f"({remote_commit_id[:8]})"
            )
    else:
        print(
            f"Fetching {sanitize_display(remote)}/{sanitize_display(target_branch)} …",
            file=sys.stderr,
        )

        if dry_run:
            # In dry-run mode we don't fetch objects — just report what would arrive.
            # We estimate commits_received from the remote pack (free from /refs response).
            if fmt == "json":
                print(json.dumps(_PullJson(
                    status="dry_run",
                    remote=remote,
                    branch=target_branch,
                    local_branch=current_branch,
                    commits_received=0,
                    objects_written=0,
                    head=None,
                    conflict_paths=[],
                    dry_run=True,
                )))
            else:
                print(
                    f"Would fetch {sanitize_display(remote)}/{sanitize_display(target_branch)} "
                    f"→ {remote_commit_id[:8]} (dry run)"
                )
            return

        # Optimised have-list: try just the local HEAD first (common case of
        # being 1–few commits behind), then fall back to full history.
        local_head = get_head_commit_id(root, current_branch)
        have_for_fetch: list[str]
        try:
            quick_have = [local_head] if local_head else []
            resp = transport.negotiate(url, signing, want=[remote_commit_id], have=quick_have)
            if resp["ready"]:
                have_for_fetch = resp["ack"] or quick_have
            else:
                all_local = [c.commit_id for c in get_all_commits(root)]
                have_for_fetch = negotiate_have(
                    transport, url, signing, [remote_commit_id], all_local
                )
        except TransportError:
            logger.debug("negotiate not supported, falling back to full have list")
            have_for_fetch = [c.commit_id for c in get_all_commits(root)]

        try:
            bundle = transport.fetch_pack(
                url, signing, want=[remote_commit_id], have=have_for_fetch
            )
        except TransportError as exc:
            print(
                f"❌ Fetch failed: {sanitize_display(str(exc))}",
                file=sys.stderr,
            )
            raise SystemExit(ExitCode.INTERNAL_ERROR)

        apply_result = apply_pack(root, bundle)
        commits_received = apply_result["commits_written"]
        objects_written = apply_result["objects_written"]
        set_remote_head(remote, target_branch, remote_commit_id, root)

        if fmt == "text":
            print(
                f"✅ Fetched {commits_received} commit(s), {objects_written} new object(s) "
                f"from {sanitize_display(remote)}/{sanitize_display(target_branch)} "
                f"({remote_commit_id[:8]})"
            )

    # ── Phase 2: ensure objects for the target snapshot ─────────────────────────
    # fetch_pack returns only VCS metadata (commits + the tip snapshot with its
    # manifest).  Before we can apply_manifest to restore the working tree, we must
    # ensure all objects listed in the target snapshot are present locally.
    # This runs even when "already_known == remote_commit_id" because the local
    # object store may be incomplete (e.g. after a failed previous clone/pull).
    def _ensure_snapshot_objects(snap_manifest: Manifest) -> int:
        """Fetch any missing objects for *snap_manifest* and return count written."""
        missing_oids = [oid for oid in snap_manifest.values() if not has_object(root, oid)]
        if not missing_oids:
            return 0
        logger.debug("pull: fetching %d missing object(s) via /fetch/objects", len(missing_oids))
        try:
            fetched_objs = transport.fetch_objects(url, signing, missing_oids)
        except TransportError as exc:
            print(
                f"❌ Fetch objects failed: {sanitize_display(str(exc))}",
                file=sys.stderr,
            )
            raise SystemExit(ExitCode.INTERNAL_ERROR)
        written = 0
        for obj in fetched_objs:
            oid = obj.get("object_id", "")
            raw = obj.get("content", b"")
            if oid and isinstance(raw, bytes) and raw:
                if write_object(root, oid, raw):
                    written += 1
        return written

    if no_merge:
        if fmt == "json":
            print(json.dumps(_PullJson(
                status="fetched",
                remote=remote,
                branch=target_branch,
                local_branch=current_branch,
                commits_received=commits_received,
                objects_written=objects_written,
                head=remote_commit_id,
                conflict_paths=[],
                dry_run=dry_run,
            )))
        return

    # ── Phase 3: merge ────────────────────────────────────────────────────────
    repo_id = read_repo_id(root)
    ours_commit_id = get_head_commit_id(root, current_branch)
    theirs_commit_id = remote_commit_id

    if ours_commit_id is None:
        # No local commits yet — bootstrap: advance HEAD to the remote commit.
        # Apply the manifest BEFORE writing the branch ref so a crash between
        # the two leaves the ref consistent with the working tree.
        if not dry_run:
            theirs_commit = read_commit(root, theirs_commit_id)
            if not theirs_commit:
                print(
                    f"❌ Pull aborted: commit {theirs_commit_id[:8]} was fetched but "
                    "is not readable from the local store (corrupt or hash mismatch). "
                    "Run `muse verify-pack` to audit the store.",
                    file=sys.stderr,
                )
                raise SystemExit(ExitCode.INTERNAL_ERROR)
            snap = read_snapshot(root, theirs_commit.snapshot_id)
            if snap is None:
                print(
                    f"❌ Pull aborted: snapshot {theirs_commit.snapshot_id[:8]} "
                    f"referenced by commit {theirs_commit_id[:8]} is missing or corrupt. "
                    "Run `muse verify-pack` to audit the store.",
                    file=sys.stderr,
                )
                raise SystemExit(ExitCode.INTERNAL_ERROR)
            objects_written += _ensure_snapshot_objects(snap.manifest)
            apply_manifest(root, snap.manifest)
            write_branch_ref(root, current_branch, theirs_commit_id)
        if fmt == "json":
            print(json.dumps(_PullJson(
                status="fast_forward",
                remote=remote,
                branch=target_branch,
                local_branch=current_branch,
                commits_received=commits_received,
                objects_written=objects_written,
                head=None if dry_run else theirs_commit_id,
                conflict_paths=[],
                dry_run=dry_run,
            )))
        else:
            suffix = " (dry run)" if dry_run else ""
            print(
                f"✅ Initialised {sanitize_display(current_branch)} "
                f"at {theirs_commit_id[:8]}{suffix}"
            )
        return

    if ours_commit_id == theirs_commit_id:
        if fmt == "json":
            print(json.dumps(_PullJson(
                status="up_to_date",
                remote=remote,
                branch=target_branch,
                local_branch=current_branch,
                commits_received=commits_received,
                objects_written=objects_written,
                head=ours_commit_id,
                conflict_paths=[],
                dry_run=dry_run,
            )))
        else:
            print("Already up to date.", file=sys.stderr)
        return

    base_commit_id = find_merge_base(root, ours_commit_id, theirs_commit_id)

    if base_commit_id == theirs_commit_id:
        # Local is already ahead of remote — nothing to pull.
        if fmt == "json":
            print(json.dumps(_PullJson(
                status="up_to_date",
                remote=remote,
                branch=target_branch,
                local_branch=current_branch,
                commits_received=commits_received,
                objects_written=objects_written,
                head=ours_commit_id,
                conflict_paths=[],
                dry_run=dry_run,
            )))
        else:
            print("Already up to date.", file=sys.stderr)
        return

    # Fast-forward: remote is a direct descendant of local HEAD.
    if base_commit_id == ours_commit_id:
        if not dry_run:
            theirs_commit = read_commit(root, theirs_commit_id)
            if not theirs_commit:
                print(
                    f"❌ Pull aborted: commit {theirs_commit_id[:8]} was fetched but "
                    "is not readable from the local store (corrupt or hash mismatch). "
                    "Run `muse verify-pack` to audit the store.",
                    file=sys.stderr,
                )
                raise SystemExit(ExitCode.INTERNAL_ERROR)
            snap = read_snapshot(root, theirs_commit.snapshot_id)
            if snap is None:
                print(
                    f"❌ Pull aborted: snapshot {theirs_commit.snapshot_id[:8]} "
                    f"referenced by commit {theirs_commit_id[:8]} is missing or corrupt. "
                    "Run `muse verify-pack` to audit the store.",
                    file=sys.stderr,
                )
                raise SystemExit(ExitCode.INTERNAL_ERROR)
            # Apply manifest BEFORE advancing the branch pointer so a
            # crash between the two leaves the ref consistent with the tree.
            objects_written += _ensure_snapshot_objects(snap.manifest)
            apply_manifest(root, snap.manifest)
            write_branch_ref(root, current_branch, theirs_commit_id)
        if fmt == "json":
            print(json.dumps(_PullJson(
                status="fast_forward",
                remote=remote,
                branch=target_branch,
                local_branch=current_branch,
                commits_received=commits_received,
                objects_written=objects_written,
                head=None if dry_run else theirs_commit_id,
                conflict_paths=[],
                dry_run=dry_run,
            )))
        else:
            suffix = " (dry run)" if dry_run else ""
            print(
                f"Fast-forward {sanitize_display(current_branch)} to "
                f"{theirs_commit_id[:8]} "
                f"({sanitize_display(remote)}/{sanitize_display(target_branch)}){suffix}"
            )
        return

    # Branches have diverged — three-way merge required.
    if ff_only:
        print(
            f"❌ Pull aborted: {sanitize_display(remote)}/{sanitize_display(target_branch)} "
            f"has diverged from {sanitize_display(current_branch)}.\n"
            f"   Fast-forward not possible. Remove --ff-only to allow a merge commit.",
            file=sys.stderr,
        )
        raise SystemExit(ExitCode.USER_ERROR)

    domain = read_domain(root)
    plugin = resolve_plugin(root)

    ours_manifest = get_head_snapshot_manifest(root, repo_id, current_branch) or {}
    theirs_commit = read_commit(root, theirs_commit_id)
    if not theirs_commit:
        print(
            f"❌ Pull aborted: commit {theirs_commit_id[:8]} is not readable from "
            "the local store (corrupt or hash mismatch). "
            "Run `muse verify-pack` to audit the store.",
            file=sys.stderr,
        )
        raise SystemExit(ExitCode.INTERNAL_ERROR)
    theirs_snap = read_snapshot(root, theirs_commit.snapshot_id)
    if theirs_snap is None:
        print(
            f"❌ Pull aborted: snapshot {theirs_commit.snapshot_id[:8]} "
            f"referenced by commit {theirs_commit_id[:8]} is missing or corrupt. "
            "A merge with a missing remote snapshot would treat all remote files "
            "as deleted. Run `muse verify-pack` to audit the store.",
            file=sys.stderr,
        )
        raise SystemExit(ExitCode.INTERNAL_ERROR)
    theirs_manifest = dict(theirs_snap.manifest)

    base_manifest: Manifest = {}
    if base_commit_id:
        base_commit = read_commit(root, base_commit_id)
        if base_commit:
            base_snap = read_snapshot(root, base_commit.snapshot_id)
            if base_snap:
                base_manifest = dict(base_snap.manifest)

    base_snap_obj = SnapshotManifest(files=base_manifest, domain=domain, directories=directories_from_manifest(base_manifest))
    ours_snap_obj = SnapshotManifest(files=ours_manifest, domain=domain, directories=directories_from_manifest(ours_manifest))
    theirs_snap_obj = SnapshotManifest(files=theirs_manifest, domain=domain, directories=directories_from_manifest(theirs_manifest))

    if isinstance(plugin, StructuredMergePlugin):
        ours_delta = plugin.diff(base_snap_obj, ours_snap_obj, repo_root=root)
        theirs_delta = plugin.diff(base_snap_obj, theirs_snap_obj, repo_root=root)
        result = plugin.merge_ops(
            base_snap_obj,
            ours_snap_obj,
            theirs_snap_obj,
            ours_delta["ops"],
            theirs_delta["ops"],
            repo_root=root,
        )
    else:
        result = plugin.merge(base_snap_obj, ours_snap_obj, theirs_snap_obj, repo_root=root)

    if fmt == "text" and result.applied_strategies:
        for p, strategy in sorted(result.applied_strategies.items()):
            if strategy != "manual":
                print(f"  ✔ [{strategy}] {p}", file=sys.stderr)

    if not result.is_clean:
        if not dry_run:
            write_merge_state(
                root,
                base_commit=base_commit_id or "",
                ours_commit=ours_commit_id,
                theirs_commit=theirs_commit_id,
                conflict_paths=result.conflicts,
                other_branch=f"{remote}/{target_branch}",
            )
        conflict_paths = sorted(result.conflicts)
        if fmt == "json":
            print(json.dumps(_PullJson(
                status="conflict",
                remote=remote,
                branch=target_branch,
                local_branch=current_branch,
                commits_received=commits_received,
                objects_written=objects_written,
                head=None,
                conflict_paths=conflict_paths,
                dry_run=dry_run,
            )))
        else:
            suffix = " (dry run — no state written)" if dry_run else ""
            print(f"❌ Merge conflict in {len(conflict_paths)} file(s){suffix}:", file=sys.stderr)
            for p in conflict_paths:
                print(f"  CONFLICT (both modified): {sanitize_display(p)}", file=sys.stderr)
            if not dry_run:
                print('\nFix conflicts and run "muse commit" to complete the merge.', file=sys.stderr)
        raise SystemExit(ExitCode.USER_ERROR)

    merged_manifest = result.merged["files"]

    if dry_run:
        if fmt == "json":
            print(json.dumps(_PullJson(
                status="dry_run",
                remote=remote,
                branch=target_branch,
                local_branch=current_branch,
                commits_received=commits_received,
                objects_written=objects_written,
                head=None,
                conflict_paths=[],
                dry_run=True,
            )))
        else:
            print(
                f"Would merge {sanitize_display(remote)}/{sanitize_display(target_branch)} "
                f"into {sanitize_display(current_branch)} (dry run)"
            )
        return

    objects_written += _ensure_snapshot_objects(merged_manifest)
    apply_manifest(root, merged_manifest)

    merged_dirs = directories_from_manifest(merged_manifest)
    snapshot_id = compute_snapshot_id(merged_manifest, merged_dirs)
    committed_at = datetime.datetime.now(datetime.timezone.utc)
    merge_message = (
        message
        or f"Merge {remote}/{target_branch} into {current_branch}"
    )
    commit_id = compute_commit_id(
        parent_ids=[ours_commit_id, theirs_commit_id],
        snapshot_id=snapshot_id,
        message=merge_message,
        committed_at_iso=committed_at.isoformat(),
    )
    write_snapshot(root, SnapshotRecord(snapshot_id=snapshot_id, manifest=merged_manifest, directories=merged_dirs))
    write_commit(
        root,
        CommitRecord(
            commit_id=commit_id,
            repo_id=repo_id,
            branch=current_branch,
            snapshot_id=snapshot_id,
            message=merge_message,
            committed_at=committed_at,
            parent_commit_id=ours_commit_id,
            parent2_commit_id=theirs_commit_id,
        ),
    )
    write_branch_ref(root, current_branch, commit_id)

    if fmt == "json":
        print(json.dumps(_PullJson(
            status="merged",
            remote=remote,
            branch=target_branch,
            local_branch=current_branch,
            commits_received=commits_received,
            objects_written=objects_written,
            head=commit_id,
            conflict_paths=[],
            dry_run=False,
        )))
    else:
        print(
            f"✅ Merged {sanitize_display(remote)}/{sanitize_display(target_branch)} "
            f"into {sanitize_display(current_branch)} ({commit_id[:8]})"
        )