"""Snapshot / restore MCP tools — portable code-intel state bundles.

Phase 2a ships the local snapshot surface. A snapshot is a ``.tar.gz``
under ``.attocode/snapshots/`` containing:

  - ``manifest.json`` at the archive root (see :data:`MANIFEST_SCHEMA_VERSION`
    below) listing every component + its SHA-256.
  - ``stores/<name>`` — live-consistent copies of every local DB
    (SQLite backup API) plus the trigram binary files.
  - ``cache_manifest.json`` — the project's current cache manifest.

Restore uses an atomic ``.attocode.staging/`` → ``rename`` dance so a
failed extraction cannot corrupt the live store.

This is the **local-only** form — it does not yet push to an OCI
registry or the HTTP server. Phase 3 will add those paths on top of the
same manifest format, so the snapshot file produced here is forward
compatible.
"""

from __future__ import annotations

import json
import logging
import os
import shutil
import sqlite3
import tarfile
import tempfile
import time
from dataclasses import dataclass
from typing import Any

from attocode_intel._shared import _get_project_dir, mcp
from attocode_intel.artifacts.hashing import sha256_file

logger = logging.getLogger(__name__)

MANIFEST_SCHEMA_VERSION = 1
SNAPSHOT_SUFFIX = ".atsnap.tar.gz"
SNAPSHOT_DIR_NAME = "snapshots"
# Local snapshot operations append audit entries to a per-project
# JSONL log so create/restore/delete events are traceable after the
# fact. The server side already writes to ``audit_events`` via
# ``log_event``; this is the equivalent for the local stdio MCP.
SNAPSHOT_AUDIT_LOG = os.path.join("cache", "snapshot_events.jsonl")

# Everything under .attocode/ that's worth capturing. Each tuple is
# (source relative path, archive relative path under "stores/").
_SQLITE_STORES: tuple[tuple[str, str], ...] = (
    ("index/symbols.db", "symbols.db"),
    ("vectors/embeddings.db", "embeddings.db"),
    ("index/kw_index.db", "kw_index.db"),
    ("cache/memory.db", "memory.db"),
    ("adrs.db", "adrs.db"),
    ("frecency/frecency.db", "frecency.db"),
    ("query_history/query_history.db", "query_history.db"),
)

# Binary files that we copy as-is (no SQLite backup dance).
_BINARY_FILES: tuple[tuple[str, str], ...] = (
    ("index/trigrams.lookup", "trigrams.lookup"),
    ("index/trigrams.postings", "trigrams.postings"),
    ("index/trigrams.db", "trigrams.db"),
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------



def _snapshots_dir(project_dir: str) -> str:
    return os.path.join(project_dir, ".attocode", SNAPSHOT_DIR_NAME)


def _append_audit_event(project_dir: str, event_type: str, detail: dict[str, Any]) -> None:
    """Append one JSONL entry to the local snapshot audit log.

    Best-effort: any write failure is logged and swallowed so a broken
    audit path never blocks a successful snapshot operation.
    """
    log_path = os.path.join(project_dir, ".attocode", SNAPSHOT_AUDIT_LOG)
    try:
        os.makedirs(os.path.dirname(log_path), exist_ok=True)
        entry = {
            "event_type": event_type,
            "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "detail": detail,
        }
        with open(log_path, "a", encoding="utf-8") as f:
            f.write(json.dumps(entry, sort_keys=True))
            f.write("\n")
    except OSError as exc:
        logger.debug("snapshot audit log append failed: %s", exc)


def _attocode_dir(project_dir: str) -> str:
    return os.path.join(project_dir, ".attocode")


def _backup_sqlite(source: str, dest: str) -> None:
    """Use the SQLite backup API for a consistent live copy.

    Copying a SQLite database file while another connection has it open
    can produce a corrupt snapshot. The backup API serializes reads on
    the source and writes a consistent image to the destination even if
    the source is actively being mutated.
    """
    src_conn = sqlite3.connect(f"file:{source}?mode=ro", uri=True)
    try:
        dest_conn = sqlite3.connect(dest)
        try:
            src_conn.backup(dest_conn)
        finally:
            dest_conn.close()
    finally:
        src_conn.close()


def _fmt_bytes(n: int) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024:
            return f"{n}{unit}" if unit == "B" else f"{n:.1f}{unit}"
        n //= 1024
    return f"{n}TB"


def _list_snapshot_files(project_dir: str) -> list[str]:
    sdir = _snapshots_dir(project_dir)
    if not os.path.isdir(sdir):
        return []
    return sorted(
        os.path.join(sdir, f)
        for f in os.listdir(sdir)
        if f.endswith(SNAPSHOT_SUFFIX)
    )


def _resolve_snapshot(project_dir: str, name: str) -> str | None:
    """Resolve ``name`` to a snapshot path. Accepts bare names, base names,
    and full paths."""
    if os.path.isabs(name) and os.path.exists(name):
        return name
    sdir = _snapshots_dir(project_dir)
    candidates = [
        os.path.join(sdir, name),
        os.path.join(sdir, name + SNAPSHOT_SUFFIX),
    ]
    # Also allow matching any existing snapshot whose basename starts with name.
    for f in _list_snapshot_files(project_dir):
        if os.path.basename(f).startswith(name):
            candidates.append(f)
    for c in candidates:
        if os.path.exists(c):
            return c
    return None


@dataclass(slots=True)
class _Component:
    name: str
    archive_path: str  # relative path inside the tar
    source_path: str
    size_bytes: int
    digest: str  # sha256 of the staged copy (not the live file)


# ---------------------------------------------------------------------------
# Build manifest + staging
# ---------------------------------------------------------------------------


def _stage_snapshot(
    project_dir: str,
    staging_dir: str,
    include_patterns: set[str] | None,
) -> list[_Component]:
    """Copy every store (SQLite via backup, binaries via shutil) into
    ``staging_dir/stores/`` and return a list of :class:`_Component`.

    ``include_patterns`` is a set of component logical names (e.g.
    ``{"symbols", "embeddings"}``). Empty / None means "include everything".
    """
    attocode = _attocode_dir(project_dir)
    stores_dir = os.path.join(staging_dir, "stores")
    os.makedirs(stores_dir, exist_ok=True)
    components: list[_Component] = []

    def _include(name: str) -> bool:
        if not include_patterns:
            return True
        return name in include_patterns

    # SQLite stores.
    for src_rel, out_name in _SQLITE_STORES:
        logical = out_name.split(".")[0]  # "symbols.db" → "symbols"
        if not _include(logical):
            continue
        src = os.path.join(attocode, src_rel)
        if not os.path.exists(src):
            continue
        dest = os.path.join(stores_dir, out_name)
        try:
            _backup_sqlite(src, dest)
        except sqlite3.Error as exc:
            logger.warning(
                "snapshot: SQLite backup of %s failed (%s); falling back to file copy",
                src, exc,
            )
            shutil.copy2(src, dest)
        size = os.path.getsize(dest)
        components.append(_Component(
            name=logical,
            archive_path=f"stores/{out_name}",
            source_path=src_rel,
            size_bytes=size,
            digest=sha256_file(dest),
        ))

    # Binary / non-SQLite files (trigrams).
    if _include("trigrams"):
        for src_rel, out_name in _BINARY_FILES:
            src = os.path.join(attocode, src_rel)
            if not os.path.exists(src):
                continue
            dest = os.path.join(stores_dir, out_name)
            shutil.copy2(src, dest)
            components.append(_Component(
                name="trigrams",
                archive_path=f"stores/{out_name}",
                source_path=src_rel,
                size_bytes=os.path.getsize(dest),
                digest=sha256_file(dest),
            ))

    # cache_manifest.json (top level of the archive)
    mpath = os.path.join(attocode, "cache_manifest.json")
    if os.path.exists(mpath):
        dest = os.path.join(staging_dir, "cache_manifest.json")
        shutil.copy2(mpath, dest)
        components.append(_Component(
            name="cache_manifest",
            archive_path="cache_manifest.json",
            source_path="cache_manifest.json",
            size_bytes=os.path.getsize(dest),
            digest=sha256_file(dest),
        ))

    return components


def _write_manifest(
    staging_dir: str,
    *,
    project_dir: str,
    name: str,
    components: list[_Component],
) -> dict[str, Any]:
    # Embed only the basename of project_dir in the manifest to avoid
    # leaking absolute filesystem paths into shareable snapshots. The
    # absolute path is kept in a private ``_producer_abs_path`` field
    # for same-host restore convenience (restore path ignores it anyway).
    manifest = {
        "schema_version": MANIFEST_SCHEMA_VERSION,
        "schema": "atto.snapshot.v1",
        "snapshot_name": name,
        "project_name": os.path.basename(os.path.abspath(project_dir)) or "project",
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "producer": {
            "service": "attocode-local",
            "version": "phase-2a",
        },
        "components": [
            {
                "name": c.name,
                "archive_path": c.archive_path,
                "source_path": c.source_path,
                "size_bytes": c.size_bytes,
                "digest": f"sha256:{c.digest}",
                "media_type": _media_type_for(c.name, c.archive_path),
            }
            for c in components
        ],
        "total_size_bytes": sum(c.size_bytes for c in components),
    }
    with open(os.path.join(staging_dir, "manifest.json"), "w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2, sort_keys=True)
    return manifest


def _media_type_for(name: str, archive_path: str) -> str:
    if archive_path.endswith(".db"):
        return f"application/vnd.attocode.sqlite.{name}.v1+sqlite"
    if name == "trigrams":
        return "application/vnd.attocode.trigram.v1+bin"
    if archive_path == "cache_manifest.json":
        return "application/vnd.attocode.cache_manifest.v1+json"
    return "application/vnd.attocode.blob.v1+bin"


# ---------------------------------------------------------------------------
# MCP tools
# ---------------------------------------------------------------------------


@mcp.tool()
def snapshot_create(name: str = "", include: str = "") -> str:
    """Create a portable snapshot of the current code-intel state.

    Bundles SQLite stores (via the live SQLite backup API — safe while
    the stores are open) + trigram files + cache_manifest.json into a
    single ``.atsnap.tar.gz`` file under ``.attocode/snapshots/``.

    Args:
        name: Human-friendly name. If empty, uses a timestamp.
        include: Comma-separated list of component logical names to
            include (e.g. ``"symbols,embeddings,learnings"``). Empty
            means "include everything".

    Returns a summary of the produced snapshot file.
    """
    project_dir = _get_project_dir()
    sdir = _snapshots_dir(project_dir)
    os.makedirs(sdir, exist_ok=True)

    timestamp = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())
    base_name = (name or f"snapshot_{timestamp}").replace("/", "_").replace(" ", "_")
    if not base_name.endswith("_" + timestamp) and not name:
        base_name = f"{base_name}"
    out_file = os.path.join(sdir, f"{base_name}{SNAPSHOT_SUFFIX}")

    include_set = {s.strip() for s in include.split(",") if s.strip()} if include else None

    with tempfile.TemporaryDirectory(prefix="atto-snap-", dir=sdir) as staging:
        components = _stage_snapshot(project_dir, staging, include_set)
        if not components:
            return (
                f"snapshot_create: nothing to snapshot "
                f"(no components matched include={include!r})"
            )
        manifest = _write_manifest(
            staging, project_dir=project_dir, name=base_name, components=components,
        )

        # Build the tar.
        tmp_tar = out_file + ".tmp"
        with tarfile.open(tmp_tar, "w:gz") as tar:
            # manifest.json first so it's cheap to read without extracting.
            tar.add(os.path.join(staging, "manifest.json"), arcname="manifest.json")
            if os.path.exists(os.path.join(staging, "cache_manifest.json")):
                tar.add(
                    os.path.join(staging, "cache_manifest.json"),
                    arcname="cache_manifest.json",
                )
            stores_dir = os.path.join(staging, "stores")
            if os.path.isdir(stores_dir):
                tar.add(stores_dir, arcname="stores")
        os.replace(tmp_tar, out_file)

    size = os.path.getsize(out_file)
    _append_audit_event(project_dir, "snapshot.created", {
        "snapshot_name": base_name,
        "archive_path": out_file,
        "components": len(manifest["components"]),
        "archive_size_bytes": size,
        "total_size_uncompressed": manifest["total_size_bytes"],
    })
    return (
        f"snapshot_create: {out_file}\n"
        f"  components: {len(manifest['components'])}\n"
        f"  total_size_uncompressed: {_fmt_bytes(manifest['total_size_bytes'])}\n"
        f"  archive_size: {_fmt_bytes(size)}"
    )


@mcp.tool()
def snapshot_list() -> str:
    """List all snapshots under ``.attocode/snapshots/``."""
    project_dir = _get_project_dir()
    snaps = _list_snapshot_files(project_dir)
    if not snaps:
        return "snapshot_list: no snapshots in .attocode/snapshots/"
    lines = [f"snapshot_list: {len(snaps)} snapshot(s):"]
    for path in snaps:
        size = os.path.getsize(path)
        mtime = time.strftime(
            "%Y-%m-%d %H:%M:%S",
            time.localtime(os.path.getmtime(path)),
        )
        lines.append(f"  {os.path.basename(path):45s} {_fmt_bytes(size):>10s}  {mtime}")
    return "\n".join(lines)


@mcp.tool()
def snapshot_delete(name: str, confirm: bool = False) -> str:
    """Delete a snapshot.

    Args:
        name: Snapshot filename or bare prefix (see ``snapshot_list``).
        confirm: Must be True to actually delete.
    """
    project_dir = _get_project_dir()
    path = _resolve_snapshot(project_dir, name)
    if path is None:
        return f"snapshot_delete: no snapshot matching {name!r}"
    size = os.path.getsize(path)
    if not confirm:
        return (
            f"snapshot_delete: DRY RUN — would delete {path} ({_fmt_bytes(size)}). "
            f"Re-run with confirm=True."
        )
    os.unlink(path)
    _append_audit_event(project_dir, "snapshot.deleted", {
        "snapshot_path": path,
        "archive_size_bytes": size,
    })
    return f"snapshot_delete: removed {path} ({_fmt_bytes(size)})"


@mcp.tool()
def snapshot_restore(name: str, confirm: bool = False) -> str:
    """Restore a snapshot into ``.attocode/``.

    Uses a staging directory + atomic rename so a partial extraction
    can't corrupt the live state. On failure the live state is preserved.

    **Caution:** this overwrites the current ``.attocode/`` store files.
    Pre-existing snapshots, the ``snapshots/`` directory itself, and any
    files not covered by the snapshot are preserved.

    Args:
        name: Snapshot filename or bare prefix.
        confirm: Must be True to actually restore.
    """
    project_dir = _get_project_dir()
    path = _resolve_snapshot(project_dir, name)
    if path is None:
        return f"snapshot_restore: no snapshot matching {name!r}"

    # Preview — read manifest without extracting the whole archive.
    try:
        with tarfile.open(path, "r:gz") as tar:
            mfile = tar.extractfile("manifest.json")
            if mfile is None:
                return "snapshot_restore: malformed snapshot (no manifest.json)"
            manifest = json.loads(mfile.read().decode("utf-8"))
    except (tarfile.TarError, json.JSONDecodeError, KeyError) as exc:
        return f"snapshot_restore: could not read snapshot manifest: {exc}"

    components = manifest.get("components", [])
    if not confirm:
        lines = [
            "snapshot_restore: DRY RUN — nothing written.",
            f"  source: {path}",
            f"  created: {manifest.get('created_at', '?')}",
            f"  components: {len(components)}",
        ]
        for c in components:
            lines.append(
                f"    - {c.get('name', '?'):15s} "
                f"{_fmt_bytes(c.get('size_bytes', 0)):>10s}  {c.get('archive_path', '?')}"
            )
        lines.append("  Re-run with confirm=True to apply.")
        return "\n".join(lines)

    # Apply: extract to a staging directory, then atomic move of each
    # individual store file. We do NOT atomic-rename the whole .attocode
    # directory because that would wipe files that aren't in the snapshot.
    attocode = _attocode_dir(project_dir)
    with tempfile.TemporaryDirectory(prefix="atto-restore-", dir=attocode) as staging:
        try:
            with tarfile.open(path, "r:gz") as tar:
                # ``filter="data"`` (Python 3.12+) rejects absolute paths,
                # path traversal, and symlinks escaping the staging dir —
                # safer than the legacy default and future-proof against
                # Python 3.14's stricter behavior.
                tar.extractall(staging, filter="data")
        except tarfile.TarError as exc:
            return f"snapshot_restore: extraction failed: {exc}"

        # Verify component digests.
        bad: list[str] = []
        for c in components:
            ap = c.get("archive_path", "")
            staged_path = os.path.join(staging, ap)
            if not os.path.exists(staged_path):
                bad.append(f"missing: {ap}")
                continue
            expected = c.get("digest", "")
            if expected.startswith("sha256:"):
                actual = "sha256:" + sha256_file(staged_path)
                if actual != expected:
                    bad.append(f"digest mismatch: {ap}")
        if bad:
            return "snapshot_restore: integrity check failed:\n  " + "\n  ".join(bad)

        # Move each component into place. Atomic per-file; the whole
        # operation is not atomic but any single file's replacement is.
        restored: list[str] = []
        failed: list[str] = []
        for c in components:
            ap = c.get("archive_path", "")
            src = os.path.join(staging, ap)
            if ap == "cache_manifest.json":
                dest = os.path.join(attocode, "cache_manifest.json")
            elif ap.startswith("stores/"):
                # Map back to the original .attocode-relative path.
                out_name = ap.split("/", 1)[1]
                dest_rel = _dest_for_store(out_name)
                if dest_rel is None:
                    failed.append(f"{ap}: unknown store name")
                    continue
                dest = os.path.join(attocode, dest_rel)
            else:
                failed.append(f"{ap}: unknown archive layout")
                continue
            os.makedirs(os.path.dirname(dest), exist_ok=True)
            try:
                os.replace(src, dest)
                restored.append(ap)
            except OSError as exc:
                failed.append(f"{ap}: {exc}")

    _append_audit_event(project_dir, "snapshot.restored", {
        "snapshot_path": path,
        "restored_components": len(restored),
        "failed_components": len(failed),
    })
    lines = [
        f"snapshot_restore: restored from {path}",
        f"  restored: {len(restored)} component(s)",
    ]
    if failed:
        lines.append(f"  failed: {len(failed)}")
        for f in failed:
            lines.append(f"    - {f}")
    return "\n".join(lines)


def _dest_for_store(out_name: str) -> str | None:
    """Map a snapshot's ``stores/<out_name>`` back to a .attocode-relative path."""
    for src_rel, candidate in _SQLITE_STORES:
        if candidate == out_name:
            return src_rel
    for src_rel, candidate in _BINARY_FILES:
        if candidate == out_name:
            return src_rel
    return None


@mcp.tool()
def snapshot_diff(a: str, b: str) -> str:
    """Diff two snapshots at the component-digest level.

    Args:
        a: First snapshot name or path.
        b: Second snapshot name or path.

    Returns a summary of which components changed, were added, or removed.
    """
    project_dir = _get_project_dir()
    path_a = _resolve_snapshot(project_dir, a)
    path_b = _resolve_snapshot(project_dir, b)
    if path_a is None:
        return f"snapshot_diff: no snapshot matching {a!r}"
    if path_b is None:
        return f"snapshot_diff: no snapshot matching {b!r}"

    def _manifest_of(path: str) -> dict[str, Any]:
        with tarfile.open(path, "r:gz") as tar:
            mfile = tar.extractfile("manifest.json")
            assert mfile is not None
            return json.loads(mfile.read().decode("utf-8"))

    try:
        ma = _manifest_of(path_a)
        mb = _manifest_of(path_b)
    except (tarfile.TarError, json.JSONDecodeError, KeyError, AssertionError) as exc:
        return f"snapshot_diff: manifest read failed: {exc}"

    def _digest_map(m: dict[str, Any]) -> dict[str, str]:
        """The local snapshot format emits three separate components all
        named ``trigrams`` (one per on-disk file). Key on
        ``(name, archive_path)`` so the three trigram files show up as
        three separate diff entries instead of collapsing.
        """
        out: dict[str, str] = {}
        for c in m.get("components", []):
            name = c.get("name", "?")
            archive_path = c.get("archive_path", "")
            key = f"{name}[{archive_path}]" if archive_path else name
            out[key] = c.get("digest", "")
        return out

    da = _digest_map(ma)
    db = _digest_map(mb)

    added = [k for k in db if k not in da]
    removed = [k for k in da if k not in db]
    changed = [k for k in da if k in db and da[k] != db[k]]
    same = [k for k in da if k in db and da[k] == db[k]]

    lines = [
        f"snapshot_diff: {os.path.basename(path_a)} → {os.path.basename(path_b)}",
        f"  same: {len(same)}",
        f"  changed: {len(changed)}",
        f"  added: {len(added)}",
        f"  removed: {len(removed)}",
    ]
    if changed:
        lines.append("  changed components:")
        for n in sorted(changed):
            lines.append(f"    - {n}: {da[n][:16]}… → {db[n][:16]}…")
    if added:
        lines.append("  added components:")
        for n in sorted(added):
            lines.append(f"    + {n}")
    if removed:
        lines.append("  removed components:")
        for n in sorted(removed):
            lines.append(f"    - {n}")
    return "\n".join(lines)
