"""
MCAL Storage Layer

Provides persistence for unified graphs across sessions.
This enables cross-session reasoning preservation - the core value proposition.

Supports two backends:
  - json  (default): File-based JSON, one file per user. Best for single-process.
  - sqlite: WAL-mode SQLite DB. Best for multi-user / multi-worker production.

Select via MCAL(storage_backend="json"|"sqlite").
"""

import base64
import json
import logging
import os
import sqlite3
import tempfile
from pathlib import Path
from typing import Optional, Protocol, runtime_checkable
from datetime import datetime, timezone

from .unified_extractor import (
    UnifiedGraph, GraphNode, GraphEdge, NodeType, EdgeType,
)


def _utc_now() -> datetime:
    """Return current UTC time (timezone-aware)."""
    return datetime.now(timezone.utc)


logger = logging.getLogger(__name__)


# =========================================================================
# Storage Protocol (Issue #127)
# =========================================================================

@runtime_checkable
class StorageBackend(Protocol):
    """Protocol that all storage backends must satisfy."""

    base_path: Path

    def save_unified_graph(self, user_id: str, graph: UnifiedGraph) -> None: ...
    def load_unified_graph(self, user_id: str) -> Optional[UnifiedGraph]: ...
    def delete_unified_graph(self, user_id: str) -> bool: ...
    def clear_user_data(self, user_id: str) -> None: ...
    def list_users(self) -> list[str]: ...


# =========================================================================
# JSON Backend (original)
# =========================================================================

class MCALStorage:
    """
    Persistent storage for MCAL data structures (JSON backend).
    
    Stores:
    - Intent graphs per user (with versioning)
    - Decision trails per user
    - Session metadata
    
    File structure:
    ~/.mcal/
      users/
        {user_id}/
          intent_graph.json      # Legacy 3-Pillar
          decisions.json          # Legacy 3-Pillar
          unified_graph.json      # Unified Deep (Issue #25)
          sessions.json
    """
    
    def __init__(self, base_path: Optional[Path] = None):
        """
        Initialize storage.
        
        Args:
            base_path: Base directory for storage (default: ~/.mcal)
        """
        if base_path is None:
            base_path = Path.home() / ".mcal"
        
        self.base_path = Path(base_path)
        self.users_path = self.base_path / "users"
        self.users_path.mkdir(parents=True, exist_ok=True)
        logger.info(f"MCAL storage initialized at {self.base_path}")
    
    def _get_user_path(self, user_id: str) -> Path:
        """Get path for user's data directory."""
        user_path = self.users_path / user_id
        user_path.mkdir(parents=True, exist_ok=True)
        return user_path
    
    def _atomic_write_json(self, target_path: Path, data: dict) -> None:
        """
        Write JSON data atomically using write-to-temp-then-rename.
        
        Issue #95: Prevents data corruption if process crashes mid-write.
        On POSIX systems, os.replace() is atomic on the same filesystem.
        
        Args:
            target_path: Final destination path for the JSON file
            data: Dictionary to serialize as JSON
        """
        target_dir = target_path.parent
        fd, tmp_path = tempfile.mkstemp(dir=target_dir, suffix='.json.tmp')
        try:
            with os.fdopen(fd, 'w') as f:
                json.dump(data, f, indent=2)
                f.flush()
                os.fsync(f.fileno())
            os.replace(tmp_path, str(target_path))  # Atomic on same filesystem
        except BaseException:
            # Clean up temp file on any failure (including KeyboardInterrupt)
            try:
                os.unlink(tmp_path)
            except OSError:
                pass
            raise
    
    # =========================================================================
    # Unified Graph Persistence (Issue #25)
    # =========================================================================
    
    def save_unified_graph(self, user_id: str, graph: UnifiedGraph) -> None:
        """
        Save unified graph for user.
        
        Args:
            user_id: User identifier
            graph: UnifiedGraph to save
        """
        user_path = self._get_user_path(user_id)
        graph_path = user_path / "unified_graph.json"
        
        # Use existing to_dict() method
        data = graph.to_dict()
        data["_metadata"] = {
            "user_id": user_id,
            "saved_at": _utc_now().isoformat(),
            "node_count": len(graph.nodes),
            "edge_count": len(graph.edges),
            "version": "unified_deep_v1"
        }
        
        self._atomic_write_json(graph_path, data)
        
        logger.info(f"Saved unified graph for {user_id}: {len(graph.nodes)} nodes, {len(graph.edges)} edges")
    
    def load_unified_graph(self, user_id: str) -> Optional[UnifiedGraph]:
        """
        Load unified graph for user.
        
        Args:
            user_id: User identifier
            
        Returns:
            UnifiedGraph or None if not found
        """
        user_path = self._get_user_path(user_id)
        graph_path = user_path / "unified_graph.json"
        
        if not graph_path.exists():
            logger.debug(f"No existing unified graph for {user_id}")
            return None
        
        try:
            with open(graph_path, 'r') as f:
                data = json.load(f)
            
            # Use existing from_dict() method
            graph = UnifiedGraph.from_dict(data)
            logger.info(f"Loaded unified graph for {user_id}: {len(graph.nodes)} nodes")
            return graph
            
        except Exception as e:
            logger.error(f"Failed to load unified graph for {user_id}: {e}")
            return None
    
    def delete_unified_graph(self, user_id: str) -> bool:
        """
        Delete unified graph for user.
        
        Args:
            user_id: User identifier
            
        Returns:
            True if deleted, False if not found
        """
        user_path = self._get_user_path(user_id)
        graph_path = user_path / "unified_graph.json"
        
        if graph_path.exists():
            graph_path.unlink()
            logger.info(f"Deleted unified graph for {user_id}")
            return True
        return False
    
    # =========================================================================
    # Utility Methods
    # =========================================================================
    
    def clear_user_data(self, user_id: str) -> None:
        """Clear all stored data for a user."""
        import shutil
        user_path = self._get_user_path(user_id)
        if user_path.exists():
            shutil.rmtree(user_path)
            logger.info(f"Cleared all data for {user_id}")
    
    def list_users(self) -> list[str]:
        """List all users with stored data."""
        if not self.users_path.exists():
            return []
        return [d.name for d in self.users_path.iterdir() if d.is_dir()]


# =========================================================================
# SQLite Backend (Issue #127)
# =========================================================================

_SQLITE_SCHEMA = """
CREATE TABLE IF NOT EXISTS graph_metadata (
    user_id     TEXT NOT NULL,
    key         TEXT NOT NULL,
    value       TEXT,
    PRIMARY KEY (user_id, key)
);

CREATE TABLE IF NOT EXISTS nodes (
    user_id     TEXT    NOT NULL,
    node_id     TEXT    NOT NULL,
    type        TEXT    NOT NULL,
    label       TEXT    NOT NULL,
    attrs_json  TEXT    NOT NULL DEFAULT '{}',
    embedding   BLOB,
    PRIMARY KEY (user_id, node_id)
);

CREATE TABLE IF NOT EXISTS edges (
    user_id     TEXT NOT NULL,
    src         TEXT NOT NULL,
    dst         TEXT NOT NULL,
    rel         TEXT NOT NULL,
    attrs_json  TEXT NOT NULL DEFAULT '{}',
    PRIMARY KEY (user_id, src, dst, rel)
);

CREATE INDEX IF NOT EXISTS idx_nodes_user   ON nodes(user_id);
CREATE INDEX IF NOT EXISTS idx_nodes_type   ON nodes(user_id, type);
CREATE INDEX IF NOT EXISTS idx_edges_user   ON edges(user_id);
CREATE INDEX IF NOT EXISTS idx_edges_src    ON edges(user_id, src);
CREATE INDEX IF NOT EXISTS idx_edges_dst    ON edges(user_id, dst);
"""


class SQLiteStorage:
    """
    SQLite storage backend for MCAL (Issue #127).

    Uses WAL mode for concurrent reads during writes.
    Single DB file at ``{base_path}/mcal.db``.

    Schema:
      - **nodes** — one row per (user_id, node_id)
      - **edges** — one row per (user_id, src, dst, rel)
      - **graph_metadata** — per-user key/value pairs (dedup config, etc.)
    """

    def __init__(self, base_path: Optional[Path] = None):
        if base_path is None:
            base_path = Path.home() / ".mcal"
        self.base_path = Path(base_path)
        self.base_path.mkdir(parents=True, exist_ok=True)

        self._db_path = self.base_path / "mcal.db"
        self._conn = sqlite3.connect(str(self._db_path), check_same_thread=False)
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._conn.execute("PRAGMA synchronous=NORMAL")
        self._conn.execute("PRAGMA foreign_keys=ON")
        self._conn.executescript(_SQLITE_SCHEMA)
        self._conn.commit()
        logger.info(f"SQLite storage initialized at {self._db_path}")

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _save_metadata(self, user_id: str, graph: UnifiedGraph) -> None:
        """Persist dedup config and graph metadata as key-value rows."""
        pairs = {
            "saved_at": _utc_now().isoformat(),
            "version": "unified_deep_v1",
            "dedup_strategy": graph._dedup_strategy.value,
            "dedup_threshold": str(graph._similarity_threshold),
        }
        if graph._type_thresholds:
            pairs["type_thresholds"] = json.dumps(
                {t.value: v for t, v in graph._type_thresholds.items()}
            )
        if graph.metadata:
            pairs["graph_metadata"] = json.dumps(graph.metadata)

        self._conn.executemany(
            "INSERT OR REPLACE INTO graph_metadata (user_id, key, value) VALUES (?, ?, ?)",
            [(user_id, k, v) for k, v in pairs.items()],
        )

    def _load_metadata(self, user_id: str) -> dict[str, str]:
        """Load key-value metadata for a user."""
        rows = self._conn.execute(
            "SELECT key, value FROM graph_metadata WHERE user_id = ?", (user_id,)
        ).fetchall()
        return {k: v for k, v in rows}

    # ------------------------------------------------------------------
    # Public API — matches MCALStorage interface
    # ------------------------------------------------------------------

    def save_unified_graph(self, user_id: str, graph: UnifiedGraph) -> None:
        """Save unified graph for user into SQLite (full replace)."""
        try:
            self._conn.execute("BEGIN")

            # 1. Delete old data for user
            self._conn.execute("DELETE FROM nodes WHERE user_id = ?", (user_id,))
            self._conn.execute("DELETE FROM edges WHERE user_id = ?", (user_id,))
            self._conn.execute("DELETE FROM graph_metadata WHERE user_id = ?", (user_id,))

            # 2. Insert nodes
            node_rows = []
            for n in graph.nodes.values():
                node_rows.append((
                    user_id,
                    n.id,
                    n.type.value,
                    n.label,
                    json.dumps(n.attrs) if n.attrs else "{}",
                    n.embedding,  # bytes or None — stored as BLOB
                ))
            if node_rows:
                self._conn.executemany(
                    "INSERT INTO nodes (user_id, node_id, type, label, attrs_json, embedding) "
                    "VALUES (?, ?, ?, ?, ?, ?)",
                    node_rows,
                )

            # 3. Insert edges
            edge_rows = []
            seen_edges: set[tuple[str, str, str, str]] = set()
            for e in graph.edges:
                key = (user_id, e.src, e.dst, e.rel.value)
                if key in seen_edges:
                    continue  # skip duplicate triples
                seen_edges.add(key)
                edge_rows.append((
                    user_id,
                    e.src,
                    e.dst,
                    e.rel.value,
                    json.dumps(e.attrs) if e.attrs else "{}",
                ))
            if edge_rows:
                self._conn.executemany(
                    "INSERT INTO edges (user_id, src, dst, rel, attrs_json) "
                    "VALUES (?, ?, ?, ?, ?)",
                    edge_rows,
                )

            # 4. Save metadata
            self._save_metadata(user_id, graph)

            self._conn.commit()
            logger.info(
                f"SQLite: saved graph for {user_id}: "
                f"{len(graph.nodes)} nodes, {len(graph.edges)} edges"
            )

        except Exception:
            self._conn.rollback()
            raise

    def load_unified_graph(self, user_id: str) -> Optional[UnifiedGraph]:
        """Load full unified graph for user from SQLite."""
        # Quick existence check
        row = self._conn.execute(
            "SELECT 1 FROM nodes WHERE user_id = ? LIMIT 1", (user_id,)
        ).fetchone()
        if row is None:
            logger.debug(f"SQLite: no graph for {user_id}")
            return None

        try:
            graph = UnifiedGraph()

            # 1. Load nodes
            for r in self._conn.execute(
                "SELECT node_id, type, label, attrs_json, embedding "
                "FROM nodes WHERE user_id = ?",
                (user_id,),
            ):
                node_id, ntype, label, attrs_json, embedding = r
                graph.nodes[node_id] = GraphNode(
                    id=node_id,
                    type=NodeType(ntype),
                    label=label,
                    attrs=json.loads(attrs_json) if attrs_json else {},
                    embedding=bytes(embedding) if embedding else None,
                )

            # 2. Load edges
            for r in self._conn.execute(
                "SELECT src, dst, rel, attrs_json FROM edges WHERE user_id = ?",
                (user_id,),
            ):
                src, dst, rel, attrs_json = r
                graph.edges.append(GraphEdge(
                    src=src,
                    dst=dst,
                    rel=EdgeType(rel),
                    attrs=json.loads(attrs_json) if attrs_json else {},
                ))

            # 3. Load metadata
            meta = self._load_metadata(user_id)
            if "dedup_strategy" in meta:
                from .unified_extractor import DeduplicationStrategy
                graph._dedup_strategy = DeduplicationStrategy(meta["dedup_strategy"])
            if "dedup_threshold" in meta:
                graph._similarity_threshold = float(meta["dedup_threshold"])
            if "type_thresholds" in meta:
                raw = json.loads(meta["type_thresholds"])
                graph._type_thresholds = {NodeType(k): v for k, v in raw.items()}
            if "graph_metadata" in meta:
                graph.metadata = json.loads(meta["graph_metadata"])

            graph._rebuild_label_index()
            logger.info(f"SQLite: loaded graph for {user_id}: {len(graph.nodes)} nodes")
            return graph

        except Exception as e:
            logger.error(f"SQLite: failed to load graph for {user_id}: {e}")
            return None

    def delete_unified_graph(self, user_id: str) -> bool:
        """Delete graph for user."""
        cur = self._conn.execute(
            "SELECT 1 FROM nodes WHERE user_id = ? LIMIT 1", (user_id,)
        )
        if cur.fetchone() is None:
            return False
        self._conn.execute("DELETE FROM nodes WHERE user_id = ?", (user_id,))
        self._conn.execute("DELETE FROM edges WHERE user_id = ?", (user_id,))
        self._conn.execute("DELETE FROM graph_metadata WHERE user_id = ?", (user_id,))
        self._conn.commit()
        logger.info(f"SQLite: deleted graph for {user_id}")
        return True

    def clear_user_data(self, user_id: str) -> None:
        """Clear all data for a user."""
        self._conn.execute("DELETE FROM nodes WHERE user_id = ?", (user_id,))
        self._conn.execute("DELETE FROM edges WHERE user_id = ?", (user_id,))
        self._conn.execute("DELETE FROM graph_metadata WHERE user_id = ?", (user_id,))
        self._conn.commit()
        logger.info(f"SQLite: cleared data for {user_id}")

    def list_users(self) -> list[str]:
        """List all users with stored data."""
        rows = self._conn.execute(
            "SELECT DISTINCT user_id FROM nodes"
        ).fetchall()
        return [r[0] for r in rows]

    def close(self) -> None:
        """Close the SQLite connection."""
        if self._conn:
            self._conn.close()
            self._conn = None

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass


# =========================================================================
# Migration Utility (Issue #127)
# =========================================================================

def migrate_json_to_sqlite(
    json_storage: MCALStorage,
    sqlite_storage: SQLiteStorage,
) -> dict[str, int]:
    """
    Migrate all user graphs from JSON backend to SQLite backend.

    Args:
        json_storage: Source JSON storage instance.
        sqlite_storage: Target SQLite storage instance.

    Returns:
        Dict with migration stats: {"users_migrated": N, "users_failed": F}.
    """
    users = json_storage.list_users()
    migrated = 0
    failed = 0

    for uid in users:
        try:
            graph = json_storage.load_unified_graph(uid)
            if graph is not None:
                sqlite_storage.save_unified_graph(uid, graph)
                migrated += 1
                logger.info(f"Migrated user {uid}: {len(graph.nodes)} nodes")
            else:
                logger.debug(f"No graph for user {uid}, skipping")
        except Exception as e:
            logger.error(f"Failed to migrate user {uid}: {e}")
            failed += 1

    result = {"users_migrated": migrated, "users_failed": failed}
    logger.info(f"Migration complete: {result}")
    return result
