"""
fw_server/persistence.py — Workspace-scoped state persistence.

Auto-save and auto-restore for both the HCP singleton and the
SessionLedger.  On shutdown the complete state is written to
FW_STATE_DIR/{hcp_state,session_state}.json.  On startup both files
are read back, restoring full context continuity across VS Code sessions.

Design constraints:
  - stdlib-only (no third-party deps)
  - Rotates up to MAX_CHECKPOINTS backups for crash safety
  - T14 (Ne Ayn Ne Gayr): checkpoint participates in the protocol's
    truth without being identical to it — a frozen projection
  - KV₇ (İhlâs): persistence is a transparent infrastructure layer;
    it never mutates or interprets HCP/session content
  - KV₃ (Observer non-interference): save/restore failures are logged,
    never propagated.  A persistence failure MUST NOT crash a tool.

Usage (via lifespan):
    state_dir is injected from FW_STATE_DIR env var at startup.
    save_hcp_state() / save_session_state() called on shutdown + atexit.
    restore_hcp_state() / restore_session_state() called on startup.
"""

from __future__ import annotations

import json
import logging
import shutil
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

logger = logging.getLogger("fw_server.persistence")

from fw_server.tracer import trace

# Maximum number of rotated checkpoint files to keep.
MAX_CHECKPOINTS = 5

# The canonical checkpoint filenames.
STATE_FILENAME = "hcp_state.json"
SESSION_FILENAME = "session_state.json"
SENTINEL_FILENAME = "sentinel_state.json"


@trace
def save_hcp_state(
    state_dir: Path,
    protocol: object,  # hcp.protocol.HCPProtocol
) -> Optional[Path]:
    """Save HCP state to *state_dir*/hcp_state.json.

    Rotates up to MAX_CHECKPOINTS previous checkpoints.
    Returns the path written, or None if nothing to save.
    """
    # Guard: nothing ingested → nothing to save
    if not getattr(protocol, "is_ready", False):
        logger.info("HCP not ready — nothing to save")
        return None

    state_dir.mkdir(parents=True, exist_ok=True)
    state_path = state_dir / STATE_FILENAME

    # Rotate existing checkpoints
    _rotate(state_path)

    # Export and write atomically (write to tmp then rename)
    try:
        state = protocol.export_state()  # type: ignore[attr-defined]
        tmp_path = state_path.with_suffix(".tmp")
        tmp_path.write_text(
            json.dumps(state, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        # Atomic rename (same filesystem)
        tmp_path.replace(state_path)
        logger.info("HCP state saved → %s (%d chunks)", state_path, len(state.get("chunks", [])))
        return state_path
    except Exception:
        logger.exception("Failed to save HCP state")
        return None


@trace
def restore_hcp_state(
    state_dir: Path,
    protocol: object,  # hcp.protocol.HCPProtocol
) -> bool:
    """Restore HCP state from *state_dir*/hcp_state.json.

    Returns True if state was restored, False otherwise.
    """
    state_path = state_dir / STATE_FILENAME
    if not state_path.exists():
        logger.info("No checkpoint found at %s — starting fresh", state_path)
        return False

    try:
        raw = state_path.read_text(encoding="utf-8")
        state = json.loads(raw)
        protocol.import_state(state)  # type: ignore[attr-defined]
        n_chunks = len(getattr(protocol, "_chunks", []))
        logger.info("HCP state restored from %s (%d chunks)", state_path, n_chunks)
        return True
    except (json.JSONDecodeError, ValueError) as e:
        logger.warning("Checkpoint corrupt or incompatible — starting fresh: %s", e)
        return False
    except Exception:
        logger.exception("Unexpected error restoring HCP state")
        return False


@trace
def get_persistence_info(state_dir: Optional[Path]) -> dict:
    """Return diagnostic info about the persistence layer.

    This is called by hcp_diagnostics to report persistence status.
    """
    info: dict = {
        "persistence_enabled": state_dir is not None,
    }
    if state_dir is None:
        info["reason"] = "FW_STATE_DIR not set"
        return info

    info["state_dir"] = str(state_dir)
    state_path = state_dir / STATE_FILENAME

    if state_path.exists():
        stat = state_path.stat()
        info["checkpoint_exists"] = True
        info["checkpoint_size_bytes"] = stat.st_size
        info["checkpoint_modified"] = datetime.fromtimestamp(
            stat.st_mtime, tz=timezone.utc,
        ).isoformat()
    else:
        info["checkpoint_exists"] = False

    # Count rotated backups
    backups = sorted(state_dir.glob("hcp_state.*.json"))
    info["n_backups"] = len(backups)

    return info


# ── Session state persistence ────────────────────────────────────────


@trace
def save_session_state(
    state_dir: Path,
    ledger: object,  # fw_server.session.SessionLedger
) -> Optional[Path]:
    """Save SessionLedger state to *state_dir*/session_state.json.

    Rotates up to MAX_CHECKPOINTS previous checkpoints.
    Returns the path written, or None if nothing to save.

    KV₃ compliance: all I/O is fire-and-forget.
    """
    # Guard: no calls recorded → nothing worth saving
    if getattr(ledger, "call_count", 0) == 0:
        logger.info("Session ledger empty — nothing to save")
        return None

    state_dir.mkdir(parents=True, exist_ok=True)
    state_path = state_dir / SESSION_FILENAME

    # Rotate existing checkpoints
    _rotate(state_path)

    try:
        state = ledger.export_state()  # type: ignore[attr-defined]
        tmp_path = state_path.with_suffix(".tmp")
        tmp_path.write_text(
            json.dumps(state, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        tmp_path.replace(state_path)
        n_calls = len(state.get("calls", []))
        logger.info(
            "Session state saved → %s (%d calls, %d score entries)",
            state_path, n_calls, len(state.get("score_ledger", [])),
        )
        return state_path
    except Exception:
        logger.exception("Failed to save session state")
        return None


@trace
def restore_session_state(
    state_dir: Path,
    ledger: object,  # fw_server.session.SessionLedger
) -> bool:
    """Restore SessionLedger state from *state_dir*/session_state.json.

    Returns True if state was restored, False otherwise.

    KV₃ compliance: corrupt/incompatible checkpoints are logged
    and silently skipped — the ledger starts fresh.
    """
    state_path = state_dir / SESSION_FILENAME
    if not state_path.exists():
        logger.info("No session checkpoint at %s — starting fresh", state_path)
        return False

    try:
        raw = state_path.read_text(encoding="utf-8")
        state = json.loads(raw)
        ledger.import_state(state)  # type: ignore[attr-defined]
        logger.info("Session state restored from %s", state_path)
        return True
    except (json.JSONDecodeError, ValueError) as e:
        logger.warning(
            "Session checkpoint corrupt or incompatible — starting fresh: %s", e,
        )
        return False
    except Exception:
        logger.exception("Unexpected error restoring session state")
        return False


@trace
def get_session_persistence_info(state_dir: Optional[Path]) -> dict:
    """Return diagnostic info about session persistence.

    Called by hcp_diagnostics to report session checkpoint status.
    """
    info: dict = {
        "session_persistence_enabled": state_dir is not None,
    }
    if state_dir is None:
        info["reason"] = "FW_STATE_DIR not set"
        return info

    info["state_dir"] = str(state_dir)
    state_path = state_dir / SESSION_FILENAME

    if state_path.exists():
        stat = state_path.stat()
        info["session_checkpoint_exists"] = True
        info["session_checkpoint_size_bytes"] = stat.st_size
        info["session_checkpoint_modified"] = datetime.fromtimestamp(
            stat.st_mtime, tz=timezone.utc,
        ).isoformat()
        # Peek at call count without loading the full ledger
        try:
            data = json.loads(state_path.read_text(encoding="utf-8"))
            info["session_call_count"] = len(data.get("calls", []))
            info["session_score_entries"] = len(data.get("score_ledger", []))
        except Exception:
            pass
    else:
        info["session_checkpoint_exists"] = False

    # Count rotated backups
    backups = sorted(state_dir.glob("session_state.*.json"))
    info["session_n_backups"] = len(backups)

    return info


# ── Sentinel state persistence ────────────────────────────────────────


@trace
def save_sentinel_state(
    state_dir: Path,
    sentinel: object,  # sentinel.engine.SentinelEngine
) -> Optional[Path]:
    """Save Sentinel state to *state_dir*/sentinel_state.json.

    Same atomic pattern as save_hcp_state / save_session_state.
    KV₃ compliance: all I/O is fire-and-forget.
    """
    if getattr(sentinel, "incident_count", 0) == 0:
        logger.info("Sentinel registry empty — nothing to save")
        return None

    state_dir.mkdir(parents=True, exist_ok=True)
    state_path = state_dir / SENTINEL_FILENAME

    _rotate(state_path)

    try:
        state = sentinel.export_state()  # type: ignore[attr-defined]
        tmp_path = state_path.with_suffix(".tmp")
        tmp_path.write_text(
            json.dumps(state, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        tmp_path.replace(state_path)
        n_incidents = len(state.get("incidents", []))
        logger.info(
            "Sentinel state saved → %s (%d incidents, %d rules)",
            state_path, n_incidents, len(state.get("pattern_rules", [])),
        )
        return state_path
    except Exception:
        logger.exception("Failed to save sentinel state")
        return None


@trace
def restore_sentinel_state(
    state_dir: Path,
    sentinel: object,  # sentinel.engine.SentinelEngine
) -> bool:
    """Restore Sentinel state from *state_dir*/sentinel_state.json.

    Returns True if state was restored, False otherwise.
    KV₃ compliance: corrupt checkpoints silently skipped.
    """
    state_path = state_dir / SENTINEL_FILENAME
    if not state_path.exists():
        logger.info("No sentinel checkpoint at %s — starting fresh", state_path)
        return False

    try:
        raw = state_path.read_text(encoding="utf-8")
        state = json.loads(raw)
        sentinel.import_state(state)  # type: ignore[attr-defined]
        logger.info("Sentinel state restored from %s", state_path)
        return True
    except (json.JSONDecodeError, ValueError) as e:
        logger.warning(
            "Sentinel checkpoint corrupt or incompatible — starting fresh: %s", e,
        )
        return False
    except Exception:
        logger.exception("Unexpected error restoring sentinel state")
        return False


@trace
def get_sentinel_persistence_info(state_dir: Optional[Path]) -> dict:
    """Return diagnostic info about sentinel persistence."""
    info: dict = {
        "sentinel_persistence_enabled": state_dir is not None,
    }
    if state_dir is None:
        info["reason"] = "FW_STATE_DIR not set"
        return info

    info["state_dir"] = str(state_dir)
    state_path = state_dir / SENTINEL_FILENAME

    if state_path.exists():
        stat = state_path.stat()
        info["sentinel_checkpoint_exists"] = True
        info["sentinel_checkpoint_size_bytes"] = stat.st_size
        info["sentinel_checkpoint_modified"] = datetime.fromtimestamp(
            stat.st_mtime, tz=timezone.utc,
        ).isoformat()
        try:
            data = json.loads(state_path.read_text(encoding="utf-8"))
            info["sentinel_incident_count"] = len(data.get("incidents", []))
            info["sentinel_pattern_rule_count"] = len(data.get("pattern_rules", []))
        except Exception:
            pass
    else:
        info["sentinel_checkpoint_exists"] = False

    backups = sorted(state_dir.glob("sentinel_state.*.json"))
    info["sentinel_n_backups"] = len(backups)

    return info


# ── Internal helpers ─────────────────────────────────────────────────

@trace
def _rotate(state_path: Path) -> None:
    """Rotate existing checkpoints, keeping up to MAX_CHECKPOINTS."""
    if not state_path.exists():
        return

    stem = state_path.stem  # e.g. "hcp_state" or "session_state"
    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
    backup = state_path.with_name(f"{stem}.{ts}.json")
    try:
        shutil.copy2(str(state_path), str(backup))
    except Exception:
        logger.warning("Failed to rotate checkpoint %s", state_path.name)
        return

    # Prune old backups
    backups = sorted(state_path.parent.glob(f"{stem}.*.json"))
    while len(backups) > MAX_CHECKPOINTS:
        oldest = backups.pop(0)
        try:
            oldest.unlink()
        except Exception:
            pass
