sync.py

"""Git Sync for DocsHaven — compressed chunks for PC <-> VPS sync.

Pattern from engram: each sync creates a NEW chunk (never modifies old ones).
No merge conflicts. Manifest tracks all chunks.
"""

import gzip
import hashlib
import io
import json
import logging
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING

from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from storage import Storage

logger = logging.getLogger(__name__)

CHUNKS_DIR = "chunks"
MANIFEST_FILE = "manifest.json"
WORK_DB = "docshaven.db"  # gitignored


class ChunkEntry(BaseModel):
    """Single chunk entry in manifest."""

    id: str  # SHA-256 prefix (16 hex chars)
    created_by: str
    created_at: str
    collections: int
    documents: int


class Manifest(BaseModel):
    """Index of all synced chunks."""

    version: int = 1
    chunks: list[ChunkEntry] = Field(default_factory=list)

    def to_dict(self) -> dict:
        return {"version": self.version, "chunks": [c.model_dump() for c in self.chunks]}

    @classmethod
    def from_dict(cls, d: dict) -> "Manifest":
        m = cls(version=d.get("version", 1))
        m.chunks = [ChunkEntry(**c) for c in d.get("chunks", [])]
        return m

    @classmethod
    def from_file(cls, path: Path) -> "Manifest":
        """Load manifest from disk, return empty if missing."""
        if not path.exists():
            return cls()
        with open(path) as f:
            return cls.from_dict(json.load(f))


_COLLECTION_PATTERN = re.compile(r"^[a-zA-Z0-9_-]+$")


class Syncer:
    """Handle compressed chunk sync between PC and VPS."""

    def __init__(self, sync_dir: Path):
        self.sync_dir = sync_dir
        self.chunks_dir = sync_dir / CHUNKS_DIR
        self.manifest_path = sync_dir / MANIFEST_FILE
        self._ensure_dirs()

    def _ensure_dirs(self) -> None:
        self.sync_dir.mkdir(parents=True, exist_ok=True)
        self.chunks_dir.mkdir(parents=True, exist_ok=True)

    def _build_export_data(self, collections_data: dict) -> dict:
        """Build chunk data structure from collections."""
        collections: dict[str, list] = {}
        total_docs = 0
        for name, docs in collections_data.items():
            collections[name] = docs
            total_docs += len(docs)
        return {
            "collections": collections,
            "exported_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        }

    def _write_chunk(self, chunk: dict, created_by: str) -> str | None:
        """Serialize, compress, and write chunk. Returns chunk_id or None if duplicate."""
        chunk_json = json.dumps(chunk, ensure_ascii=False).encode()
        chunk_id = hashlib.sha256(chunk_json).hexdigest()[:16]
        chunk_path = self.chunks_dir / f"{chunk_id}.json.gz"
        with gzip.open(chunk_path, "wb") as f:
            f.write(chunk_json)
        return chunk_id

    def export(self, collections_data: dict, created_by: str = "unknown") -> dict:
        """Export collections data as a compressed chunk.

        Args:
            collections_data: {collection_name: [documents]}
            created_by: Username or machine identifier

        Returns:
            {chunk_id, collections, documents, isEmpty}
        """
        manifest = Manifest.from_file(self.manifest_path)
        chunk = self._build_export_data(collections_data)
        total_docs = sum(len(d) for d in chunk["collections"].values())
        if total_docs == 0:
            return {"isEmpty": True}
        chunk_json = json.dumps(chunk, ensure_ascii=False).encode()
        chunk_id = hashlib.sha256(chunk_json).hexdigest()[:16]
        known = {c.id for c in manifest.chunks}
        if chunk_id in known:
            return {"isEmpty": True, "duplicate": True}
        chunk_path = self.chunks_dir / f"{chunk_id}.json.gz"
        with gzip.open(chunk_path, "wb") as f:
            f.write(chunk_json)
        entry = ChunkEntry(
            id=chunk_id,
            created_by=created_by,
            created_at=datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
            collections=len(collections_data),
            documents=total_docs,
        )
        manifest.chunks.append(entry)
        self._write_manifest(manifest)
        return {"chunk_id": chunk_id, "collections": len(collections_data), "documents": total_docs, "isEmpty": False}

    def _import_chunk_data(self, chunk_data: dict, storage) -> tuple[int, int]:
        """Import a single chunk's data into storage. Returns (collections, docs) counts."""
        collections = chunk_data.get("collections", {})
        if storage is None:
            return len(collections), sum(len(d) for d in collections.values())

        documents = []
        skipped = 0
        for collection_name, docs in collections.items():
            if not _COLLECTION_PATTERN.match(collection_name):
                logger.warning("Skipping invalid collection name: %s", collection_name)
                skipped += 1
                continue
            for doc in docs:
                if isinstance(doc, dict):
                    documents.append(
                        {
                            "collection": collection_name,
                            "path": doc.get("path", ""),
                            "content": doc.get("content", ""),
                            "title": doc.get("title", ""),
                        }
                    )
        if documents:
            storage.bulk_insert(documents)
        return len(collections) - skipped, sum(len(d) for d in collections.values())

    def import_chunks(self, storage: "Storage | None" = None) -> dict:
        """Import all chunks not yet applied.

        Args:
            storage: Optional Storage instance to import documents into.

        Returns:
            {chunks_imported, collections_imported, documents_imported, chunks_skipped}
        """
        manifest = Manifest.from_file(self.manifest_path)
        if not manifest.chunks:
            return {"chunks_imported": 0}

        result = {
            "chunks_imported": 0,
            "collections_imported": 0,
            "documents_imported": 0,
            "chunks_skipped": 0,
        }

        _HEX16 = re.compile(r"^[0-9a-f]{16}$")

        for entry in manifest.chunks:
            if not _HEX16.match(entry.id):
                logger.warning("Skipping chunk with invalid ID: %s", entry.id)
                result["chunks_skipped"] += 1
                continue
            chunk_path = self.chunks_dir / f"{entry.id}.json.gz"
            if not chunk_path.exists():
                result["chunks_skipped"] += 1
                continue

            # Guard against oversized chunks (max 10MB compressed, ~30MB uncompressed)
            if chunk_path.stat().st_size > 10 * 1024 * 1024:
                result["chunks_skipped"] += 1
                continue

            # Check if this chunk was already imported (by checking content hash)
            if storage is not None:
                existing = storage.count_documents_by_path(f"%{entry.id}%")
                if existing > 0:
                    result["chunks_skipped"] += 1
                    continue

            try:
                MAX_DECOMPRESSED = 50 * 1024 * 1024
                buf = io.BytesIO()
                with gzip.open(chunk_path, "rb") as f:
                    while chunk := f.read(8192):
                        buf.write(chunk)
                        if buf.tell() > MAX_DECOMPRESSED:
                            result["chunks_skipped"] += 1
                            break
                if buf.tell() > MAX_DECOMPRESSED:
                    continue
                chunk_data = json.loads(buf.getvalue())
            except (OSError, json.JSONDecodeError) as e:
                logger.warning("Skipping corrupted chunk %s: %s", entry.id, e)
                result["chunks_skipped"] += 1
                continue

            n_collections, n_docs = self._import_chunk_data(chunk_data, storage)
            result["chunks_imported"] += 1
            result["collections_imported"] += n_collections
            result["documents_imported"] += n_docs

        return result

    def status(self) -> dict:
        """Get sync status."""
        manifest = Manifest.from_file(self.manifest_path)
        local_chunks = len(manifest.chunks)

        # Count actual chunk files
        actual_files = len(list(self.chunks_dir.glob("*.json.gz")))

        return {
            "local_chunks": local_chunks,
            "chunk_files": actual_files,
            "manifest_size": self.manifest_path.stat().st_size if self.manifest_path.exists() else 0,
        }

    def _write_manifest(self, manifest: Manifest) -> None:
        """Atomic write: temp file + rename."""
        tmp_path = self.manifest_path.with_suffix(".tmp")
        try:
            with open(tmp_path, "w") as f:
                json.dump(manifest.to_dict(), f, indent=2)
            tmp_path.replace(self.manifest_path)
        except OSError:
            if tmp_path.exists():
                tmp_path.unlink()
            raise


def get_username() -> str:
    """Get current username for chunk attribution."""
    import getpass

    try:
        return getpass.getuser()
    except Exception:
        return "unknown"