"""
DAG Memory Layer — Phase 23.5

Memory is no longer a flat list — it forms a directed acyclic graph (DAG)
where raw observations become leaf nodes and summaries form higher-level
parent nodes.

Key invariant: Raw nodes are NEVER deleted or overwritten.
               Summaries reference children, never replace them.

Node types:
    raw      — directly ingested observation (leaf node)
    summary  — compressed summary of multiple raw nodes (parent node)
"""

import json
import time
import uuid
from dataclasses import dataclass, field
from typing import List, Optional

from openmemo.core.observation import Observation


@dataclass
class MemoryNode:
    """
    A node in the DAG Memory Graph.

    Fields:
        id          — unique node ID (node_<hex>)
        type        — "raw" or "summary"
        content     — text content of this node
        parent_ids  — parent summary nodes (empty for orphan roots)
        child_ids   — raw children this summary was built from
        agent_id    — owning agent
        scene       — memory scene/domain
        task_id     — associated task
        timestamp   — creation time
        importance  — 0.0–1.0 relevance weight
        obs_id      — original observation ID (for raw nodes)
        metadata    — arbitrary extra fields
    """
    id: str = field(default_factory=lambda: f"node_{uuid.uuid4().hex[:12]}")
    type: str = "raw"
    content: str = ""
    parent_ids: List[str] = field(default_factory=list)
    child_ids: List[str] = field(default_factory=list)
    agent_id: str = ""
    scene: str = ""
    task_id: str = ""
    timestamp: float = field(default_factory=time.time)
    importance: float = 0.5
    obs_id: str = ""
    metadata: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "type": self.type,
            "content": self.content,
            "parent_ids": self.parent_ids,
            "child_ids": self.child_ids,
            "agent_id": self.agent_id,
            "scene": self.scene,
            "task_id": self.task_id,
            "timestamp": self.timestamp,
            "importance": self.importance,
            "obs_id": self.obs_id,
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, d: dict) -> "MemoryNode":
        parent_ids = d.get("parent_ids", [])
        child_ids = d.get("child_ids", [])
        if isinstance(parent_ids, str):
            try:
                parent_ids = json.loads(parent_ids)
            except Exception:
                parent_ids = []
        if isinstance(child_ids, str):
            try:
                child_ids = json.loads(child_ids)
            except Exception:
                child_ids = []
        metadata = d.get("metadata", {})
        if isinstance(metadata, str):
            try:
                metadata = json.loads(metadata)
            except Exception:
                metadata = {}
        return cls(
            id=d.get("id", f"node_{uuid.uuid4().hex[:12]}"),
            type=d.get("type", "raw"),
            content=d.get("content", ""),
            parent_ids=parent_ids,
            child_ids=child_ids,
            agent_id=d.get("agent_id", ""),
            scene=d.get("scene", ""),
            task_id=d.get("task_id", ""),
            timestamp=float(d.get("timestamp", time.time())),
            importance=float(d.get("importance", 0.5)),
            obs_id=d.get("obs_id", ""),
            metadata=metadata,
        )

    @classmethod
    def from_observation(cls, obs: Observation) -> "MemoryNode":
        """Create a raw leaf node from an Observation."""
        return cls(
            type="raw",
            content=obs.to_text(),
            agent_id=obs.agent_id,
            scene=obs.context.get("scene", ""),
            task_id=obs.task_id,
            timestamp=obs.timestamp,
            importance=obs.confidence,
            obs_id=obs.id,
            metadata={"obs_type": obs.type, "content": obs.content,
                      "context": obs.context},
        )


class MemoryGraph:
    """
    DAG Memory Graph — stores and retrieves MemoryNodes.

    Uses the SQLiteStore's memory_nodes table.
    Provides: ingest, get_node, list_raw, list_summaries,
              get_children, get_parents, add_summary.
    """

    def __init__(self, store):
        self._store = store

    def ingest(self, obs: Observation) -> MemoryNode:
        """Convert an Observation to a raw node and persist it."""
        node = MemoryNode.from_observation(obs)
        self._store.put_memory_node(node.to_dict())
        return node

    def add_summary(self, content: str, child_ids: List[str],
                    agent_id: str = "", scene: str = "",
                    task_id: str = "", importance: float = 0.7) -> MemoryNode:
        """
        Create a summary node that references raw child nodes.
        Does NOT delete children (DAG invariant preserved).
        """
        node = MemoryNode(
            type="summary",
            content=content,
            child_ids=child_ids,
            agent_id=agent_id,
            scene=scene,
            task_id=task_id,
            importance=importance,
        )
        self._store.put_memory_node(node.to_dict())

        for cid in child_ids:
            child = self.get_node(cid)
            if child and node.id not in child.parent_ids:
                child.parent_ids.append(node.id)
                self._store.put_memory_node(child.to_dict())

        return node

    def get_node(self, node_id: str) -> Optional[MemoryNode]:
        data = self._store.get_memory_node(node_id)
        if not data:
            return None
        return MemoryNode.from_dict(data)

    def list_raw(self, agent_id: str = "", scene: str = "",
                 task_id: str = "", limit: int = 100) -> List[MemoryNode]:
        rows = self._store.list_memory_nodes(
            node_type="raw", agent_id=agent_id,
            scene=scene, task_id=task_id, limit=limit
        )
        return [MemoryNode.from_dict(r) for r in rows]

    def list_summaries(self, agent_id: str = "", scene: str = "",
                       limit: int = 50) -> List[MemoryNode]:
        rows = self._store.list_memory_nodes(
            node_type="summary", agent_id=agent_id,
            scene=scene, limit=limit
        )
        return [MemoryNode.from_dict(r) for r in rows]

    def get_children(self, node_id: str) -> List[MemoryNode]:
        """Get all raw children of a summary node."""
        node = self.get_node(node_id)
        if not node:
            return []
        results = []
        for cid in node.child_ids:
            child = self.get_node(cid)
            if child:
                results.append(child)
        return results

    def get_parents(self, node_id: str) -> List[MemoryNode]:
        """Get all summary parents of a raw node."""
        node = self.get_node(node_id)
        if not node:
            return []
        results = []
        for pid in node.parent_ids:
            parent = self.get_node(pid)
            if parent:
                results.append(parent)
        return results

    def count_uncompressed_raw(self, agent_id: str = "",
                               scene: str = "") -> int:
        """Count raw nodes that have no parent (not yet compressed)."""
        raw_nodes = self.list_raw(agent_id=agent_id, scene=scene, limit=1000)
        return sum(1 for n in raw_nodes if not n.parent_ids)
