"""
Belief Revision System - Content-Addressable Storage

Provides immutable, content-addressed storage for all BRS objects.
Supports world versioning, proposal queuing, and efficient querying.
"""
from __future__ import annotations

import datetime
import json
import sqlite3
from pathlib import Path
from typing import Any, Dict, List, Optional, Protocol, Tuple, Type, TypeVar

from .core import (
    canonical_json, content_hash, to_dict,
    WorldBundle, Node, Edge, Evidence, Pattern, Proposal,
    ObjectKind, ObjectHash, DomainId, VersionLabel
)


# =============================================================================
# Storage Backend Protocol
# =============================================================================

class StorageBackend(Protocol):
    """Protocol for storage backends (SQLite, PostgreSQL, in-memory, etc.)"""

    def put_object(self, kind: ObjectKind, obj: Any) -> ObjectHash: ...
    def get_object(self, h: ObjectHash) -> Dict[str, Any]: ...
    def put_world(self, world: WorldBundle) -> ObjectHash: ...
    def get_world(self, domain_id: DomainId, version_label: VersionLabel) -> Dict[str, Any]: ...
    def fork_world(self, domain_id: DomainId, from_version: VersionLabel, to_version: VersionLabel, notes: str = "") -> ObjectHash: ...
    def close(self) -> None: ...


# =============================================================================
# SQLite Schema
# =============================================================================

SCHEMA = """
-- Content-addressable object store
CREATE TABLE IF NOT EXISTS objects (
  hash TEXT PRIMARY KEY,
  kind TEXT NOT NULL,
  json TEXT NOT NULL
);

-- Node index for efficient queries
CREATE TABLE IF NOT EXISTS nodes (
  node_id TEXT PRIMARY KEY,
  name TEXT NOT NULL,
  hash TEXT NOT NULL,
  FOREIGN KEY(hash) REFERENCES objects(hash)
);

-- Edge index with parent/child lookups
CREATE TABLE IF NOT EXISTS edges (
  edge_id TEXT PRIMARY KEY,
  parent_id TEXT NOT NULL,
  child_id TEXT NOT NULL,
  relation TEXT NOT NULL,
  tier INTEGER NOT NULL,
  confidence REAL NOT NULL,
  hash TEXT NOT NULL,
  FOREIGN KEY(hash) REFERENCES objects(hash)
);

CREATE INDEX IF NOT EXISTS idx_edges_parent ON edges(parent_id);
CREATE INDEX IF NOT EXISTS idx_edges_child ON edges(child_id);

-- Proposal queue with status tracking
CREATE TABLE IF NOT EXISTS proposals (
  proposal_id TEXT PRIMARY KEY,
  kind TEXT NOT NULL,
  status TEXT NOT NULL,
  hash TEXT NOT NULL,
  created_utc TEXT NOT NULL,
  updated_utc TEXT NOT NULL,
  notes TEXT NOT NULL DEFAULT '',
  score REAL DEFAULT NULL,
  eval_json TEXT NOT NULL DEFAULT '',
  FOREIGN KEY(hash) REFERENCES objects(hash)
);

CREATE INDEX IF NOT EXISTS idx_proposals_status ON proposals(status);

-- World bundle registry (domain, version) -> hash
CREATE TABLE IF NOT EXISTS worlds (
  domain_id TEXT NOT NULL,
  version_label TEXT NOT NULL,
  hash TEXT NOT NULL,
  created_utc TEXT NOT NULL,
  PRIMARY KEY(domain_id, version_label),
  FOREIGN KEY(hash) REFERENCES objects(hash)
);
"""


def _now_utc() -> str:
    """Current UTC timestamp in ISO format."""
    return datetime.datetime.utcnow().isoformat() + "Z"


def _column_exists(conn: sqlite3.Connection, table: str, column: str) -> bool:
    """Check if a column exists in a table."""
    cur = conn.cursor()
    rows = cur.execute(f"PRAGMA table_info({table})").fetchall()
    return any(r[1] == column for r in rows)


def _migrate(conn: sqlite3.Connection) -> None:
    """Run lightweight schema migrations."""
    if not _column_exists(conn, 'proposals', 'score'):
        conn.execute('ALTER TABLE proposals ADD COLUMN score REAL DEFAULT NULL')
    if not _column_exists(conn, 'proposals', 'eval_json'):
        conn.execute("ALTER TABLE proposals ADD COLUMN eval_json TEXT NOT NULL DEFAULT ''")
    conn.commit()


# =============================================================================
# CAS Store Implementation
# =============================================================================

class CASStore:
    """
    Content-Addressable Storage for BRS.

    All objects are stored by their content hash, ensuring:
    - Immutability: Objects never change once stored
    - Deduplication: Identical objects share storage
    - Verifiability: Hash proves integrity
    """

    def __init__(self, root: Path, in_memory: bool = False):
        """
        Initialize storage.

        Args:
            root: Directory for storage files
            in_memory: If True, use in-memory SQLite (for testing)
        """
        self.root = root
        self.in_memory = in_memory

        if in_memory:
            self._conn = sqlite3.connect(":memory:")
        else:
            self.root.mkdir(parents=True, exist_ok=True)
            self.db_path = self.root / "brs.sqlite"
            self._conn = sqlite3.connect(str(self.db_path))
            self._conn.execute("PRAGMA journal_mode=WAL;")

        self._conn.executescript(SCHEMA)
        self._conn.commit()
        _migrate(self._conn)

    def close(self) -> None:
        """Close database connection."""
        self._conn.close()

    # -------------------------------------------------------------------------
    # Core Object Operations
    # -------------------------------------------------------------------------

    def put_object(self, kind: ObjectKind, obj: Any) -> ObjectHash:
        """
        Store an object in CAS.

        Args:
            kind: Object type (e.g., "Node", "Edge", "Pattern")
            obj: Object to store (dataclass or dict)

        Returns:
            Content hash of the stored object
        """
        payload = to_dict(obj)
        h = content_hash(payload)
        js = canonical_json(payload)
        cur = self._conn.cursor()
        cur.execute(
            "INSERT OR IGNORE INTO objects(hash, kind, json) VALUES(?,?,?)",
            (h, kind, js),
        )
        self._conn.commit()
        return h

    def get_object(self, h: ObjectHash) -> Dict[str, Any]:
        """
        Retrieve an object by hash.

        Args:
            h: Content hash

        Returns:
            Dict with 'kind' and 'json' keys

        Raises:
            KeyError: If object not found
        """
        cur = self._conn.cursor()
        row = cur.execute("SELECT kind, json FROM objects WHERE hash=?", (h,)).fetchone()
        if row is None:
            raise KeyError(f"Object not found: {h}")
        kind, js = row
        return {"kind": kind, "json": json.loads(js)}

    def object_exists(self, h: ObjectHash) -> bool:
        """Check if an object exists."""
        cur = self._conn.cursor()
        row = cur.execute("SELECT 1 FROM objects WHERE hash=?", (h,)).fetchone()
        return row is not None

    def list_objects_by_kind(self, kind: ObjectKind) -> List[Dict[str, Any]]:
        """List all objects of a given kind."""
        cur = self._conn.cursor()
        rows = cur.execute("SELECT hash, json FROM objects WHERE kind=?", (kind,)).fetchall()
        return [{"hash": r[0], "data": json.loads(r[1])} for r in rows]

    # -------------------------------------------------------------------------
    # Node Operations
    # -------------------------------------------------------------------------

    def upsert_node(self, node_id: str, name: str, obj_hash: ObjectHash) -> None:
        """Index a node for efficient lookup."""
        self._conn.execute(
            "INSERT OR REPLACE INTO nodes(node_id, name, hash) VALUES(?,?,?)",
            (node_id, name, obj_hash),
        )
        self._conn.commit()

    def get_node_by_id(self, node_id: str) -> Dict[str, Any]:
        """Get node index entry by ID."""
        cur = self._conn.cursor()
        row = cur.execute("SELECT node_id, name, hash FROM nodes WHERE node_id=?", (node_id,)).fetchone()
        if row is None:
            raise KeyError(f"Node not found: {node_id}")
        return {"node_id": row[0], "name": row[1], "hash": row[2]}

    def search_nodes_by_name(self, query: str, limit: int = 10) -> List[Dict[str, Any]]:
        """Search nodes by name substring."""
        cur = self._conn.cursor()
        like = f"%{query}%"
        rows = cur.execute(
            "SELECT node_id, name, hash FROM nodes WHERE name LIKE ? ORDER BY name LIMIT ?",
            (like, limit)
        ).fetchall()
        return [{"node_id": r[0], "name": r[1], "hash": r[2]} for r in rows]

    # -------------------------------------------------------------------------
    # Edge Operations
    # -------------------------------------------------------------------------

    def upsert_edge(
        self,
        edge_id: str,
        parent_id: str,
        child_id: str,
        relation: str,
        tier: int,
        confidence: float,
        obj_hash: ObjectHash
    ) -> None:
        """Index an edge for efficient traversal."""
        self._conn.execute(
            "INSERT OR REPLACE INTO edges(edge_id, parent_id, child_id, relation, tier, confidence, hash) VALUES(?,?,?,?,?,?,?)",
            (edge_id, parent_id, child_id, relation, tier, confidence, obj_hash),
        )
        self._conn.commit()

    def list_edges_from(self, parent_id: str) -> List[Dict[str, Any]]:
        """List all edges from a parent node (children of parent)."""
        cur = self._conn.cursor()
        rows = cur.execute(
            "SELECT edge_id, parent_id, child_id, relation, tier, confidence, hash FROM edges WHERE parent_id=?",
            (parent_id,),
        ).fetchall()
        return [
            {
                "edge_id": r[0], "parent_id": r[1], "child_id": r[2],
                "relation": r[3], "tier": r[4], "confidence": r[5], "hash": r[6]
            }
            for r in rows
        ]

    def list_edges_into(self, child_id: str) -> List[Dict[str, Any]]:
        """List all edges into a child node (parents of child)."""
        cur = self._conn.cursor()
        rows = cur.execute(
            "SELECT edge_id, parent_id, child_id, relation, tier, confidence, hash FROM edges WHERE child_id=?",
            (child_id,),
        ).fetchall()
        return [
            {
                "edge_id": r[0], "parent_id": r[1], "child_id": r[2],
                "relation": r[3], "tier": r[4], "confidence": r[5], "hash": r[6]
            }
            for r in rows
        ]

    # -------------------------------------------------------------------------
    # World Operations
    # -------------------------------------------------------------------------

    def put_world(self, world: WorldBundle) -> ObjectHash:
        """
        Store a world bundle.

        Args:
            world: WorldBundle to store

        Returns:
            Content hash of the world
        """
        h = self.put_object("WorldBundle", world)
        self._conn.execute(
            "INSERT OR REPLACE INTO worlds(domain_id, version_label, hash, created_utc) VALUES(?,?,?,?)",
            (world.domain_id, world.version_label, h, world.created_utc),
        )
        self._conn.commit()
        return h

    def get_world_hash(self, domain_id: DomainId, version_label: VersionLabel) -> ObjectHash:
        """Get the hash of a world bundle."""
        cur = self._conn.cursor()
        row = cur.execute(
            "SELECT hash FROM worlds WHERE domain_id=? AND version_label=?",
            (domain_id, version_label)
        ).fetchone()
        if row is None:
            raise KeyError(f"World not found: {domain_id}:{version_label}")
        return row[0]

    def get_world(self, domain_id: DomainId, version_label: VersionLabel) -> Dict[str, Any]:
        """Get a world bundle by domain and version."""
        h = self.get_world_hash(domain_id, version_label)
        return self.get_object(h)

    def world_exists(self, domain_id: DomainId, version_label: VersionLabel) -> bool:
        """Check if a world exists."""
        cur = self._conn.cursor()
        row = cur.execute(
            "SELECT 1 FROM worlds WHERE domain_id=? AND version_label=?",
            (domain_id, version_label)
        ).fetchone()
        return row is not None

    def list_worlds(self, domain_id: Optional[DomainId] = None) -> List[Dict[str, Any]]:
        """List all worlds, optionally filtered by domain."""
        cur = self._conn.cursor()
        if domain_id:
            rows = cur.execute(
                "SELECT domain_id, version_label, hash, created_utc FROM worlds WHERE domain_id=? ORDER BY created_utc DESC",
                (domain_id,)
            ).fetchall()
        else:
            rows = cur.execute(
                "SELECT domain_id, version_label, hash, created_utc FROM worlds ORDER BY domain_id, created_utc DESC"
            ).fetchall()
        return [
            {"domain_id": r[0], "version_label": r[1], "hash": r[2], "created_utc": r[3]}
            for r in rows
        ]

    def fork_world(
        self,
        domain_id: DomainId,
        from_version: VersionLabel,
        to_version: VersionLabel,
        notes: str = ""
    ) -> ObjectHash:
        """
        Fork a world to a new version.

        Creates a copy of the source world with a new version label.
        The original world remains unchanged.

        Args:
            domain_id: Domain identifier
            from_version: Source version label
            to_version: Target version label
            notes: Optional notes about the fork

        Returns:
            Hash of the new world
        """
        w = self.get_world(domain_id, from_version)["json"]
        w["version_label"] = to_version
        w["notes"] = notes or f"forked from {from_version}"
        w["created_utc"] = _now_utc()

        h = content_hash(w)
        js = canonical_json(w)
        self._conn.execute(
            "INSERT OR IGNORE INTO objects(hash, kind, json) VALUES(?,?,?)",
            (h, "WorldBundle", js)
        )
        self._conn.execute(
            "INSERT OR REPLACE INTO worlds(domain_id, version_label, hash, created_utc) VALUES(?,?,?,?)",
            (domain_id, to_version, h, w["created_utc"]),
        )
        self._conn.commit()
        return h

    # -------------------------------------------------------------------------
    # Proposal Operations
    # -------------------------------------------------------------------------

    def put_proposal(
        self,
        kind: str,
        proposal_obj: Any,
        status: str = "new",
        notes: str = ""
    ) -> str:
        """
        Store a proposal in the queue.

        Args:
            kind: Proposal type (e.g., "AnalogProposal")
            proposal_obj: Proposal data (dict with optional 'id' key)
            status: Initial status
            notes: Optional notes

        Returns:
            Proposal ID
        """
        h = self.put_object(kind, proposal_obj)
        pid = proposal_obj.get("id") if isinstance(proposal_obj, dict) else None
        if not pid:
            pid = h
        now = _now_utc()
        self._conn.execute(
            "INSERT OR REPLACE INTO proposals(proposal_id, kind, status, hash, created_utc, updated_utc, notes, score, eval_json) VALUES(?,?,?,?,?,?,?,?,?)",
            (pid, kind, status, h, now, now, notes or "", None, ""),
        )
        self._conn.commit()
        return pid

    def get_proposal(self, proposal_id: str) -> Dict[str, Any]:
        """Get a proposal by ID."""
        cur = self._conn.cursor()
        row = cur.execute(
            "SELECT proposal_id, kind, status, hash, created_utc, updated_utc, notes, score, eval_json FROM proposals WHERE proposal_id=?",
            (proposal_id,),
        ).fetchone()
        if row is None:
            raise KeyError(f"Proposal not found: {proposal_id}")
        obj = self.get_object(row[3])
        return {
            "proposal_id": row[0], "kind": row[1], "status": row[2],
            "hash": row[3], "created_utc": row[4], "updated_utc": row[5],
            "notes": row[6], "score": row[7], "eval_json": row[8],
            "object": obj,
        }

    def list_proposals(
        self,
        status: Optional[str] = None,
        limit: int = 50
    ) -> List[Dict[str, Any]]:
        """List proposals, optionally filtered by status."""
        cur = self._conn.cursor()
        if status:
            rows = cur.execute(
                "SELECT proposal_id, kind, status, hash, created_utc, updated_utc, notes, score, eval_json FROM proposals WHERE status=? ORDER BY (score IS NULL) ASC, score DESC, updated_utc DESC LIMIT ?",
                (status, limit),
            ).fetchall()
        else:
            rows = cur.execute(
                "SELECT proposal_id, kind, status, hash, created_utc, updated_utc, notes, score, eval_json FROM proposals ORDER BY (score IS NULL) ASC, score DESC, updated_utc DESC LIMIT ?",
                (limit,),
            ).fetchall()
        return [
            {
                "proposal_id": r[0], "kind": r[1], "status": r[2],
                "hash": r[3], "created_utc": r[4], "updated_utc": r[5],
                "notes": r[6], "score": r[7], "eval_json": r[8],
            }
            for r in rows
        ]

    def update_proposal_eval(
        self,
        proposal_id: str,
        score: float,
        eval_json: str,
        notes: str = "",
        status: Optional[str] = None
    ) -> None:
        """Update proposal with evaluation results."""
        now = _now_utc()
        if status is None:
            self._conn.execute(
                "UPDATE proposals SET score=?, eval_json=?, notes=?, updated_utc=? WHERE proposal_id=?",
                (score, eval_json or "", notes or "", now, proposal_id),
            )
        else:
            self._conn.execute(
                "UPDATE proposals SET score=?, eval_json=?, notes=?, status=?, updated_utc=? WHERE proposal_id=?",
                (score, eval_json or "", notes or "", status, now, proposal_id),
            )
        self._conn.commit()

    def set_proposal_status(self, proposal_id: str, status: str, notes: str = "") -> None:
        """Update proposal status."""
        now = _now_utc()
        self._conn.execute(
            "UPDATE proposals SET status=?, updated_utc=?, notes=? WHERE proposal_id=?",
            (status, now, notes or "", proposal_id),
        )
        self._conn.commit()

    # -------------------------------------------------------------------------
    # Bulk Operations
    # -------------------------------------------------------------------------

    def count_objects_by_kind(self) -> Dict[str, int]:
        """Count objects grouped by kind."""
        cur = self._conn.cursor()
        rows = cur.execute("SELECT kind, COUNT(*) FROM objects GROUP BY kind").fetchall()
        return {r[0]: r[1] for r in rows}

    def vacuum(self) -> None:
        """Optimize database storage."""
        self._conn.execute("VACUUM")
