"""LociDB — the main database interface.

Provides memorize(), recall(), traverse(), and related methods that
orchestrate the store, indexer, and scorer.
"""

from __future__ import annotations

import json
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from .conflict import ConflictWarning, check_conflicts as _check_conflicts, check_new_memory_conflicts
from .edge import MemoryEdge
from .health import HealthReport, compute_health, get_onboarding_questions as _get_onboarding_questions
from .indexer import Indexer
from .node import MemoryNode
from .persona import (
    clear_persona as _clear_persona,
    format_persona_prompt as _format_persona_prompt,
    get_all_persona as _get_all_persona,
    get_persona as _get_persona,
    set_persona as _set_persona,
    update_persona as _update_persona,
)
from .portability import (
    export_context_prompt,
    export_native,
    export_snapshot,
    import_context_md,
    import_native,
)
from .scorer import compute_importance
from .store import MemoryStore
from .summarizer import compress_daily
from .trace import (
    explain_trace as _explain_trace,
    get_trace as _get_trace,
    list_traces as _list_traces,
    store_trace as _store_trace,
    traces_using_node as _traces_using_node,
)


class LociDB:
    """Embedded cognitive database for LLM memory.

    Usage::

        loci = LociDB(path="./memory.loci")
        loci.memorize("Decided to use hub-spoke model", entities=["routing"])
        results = loci.recall("delivery optimization", top_k=5)
    """

    def __init__(
        self,
        path: str | Path = "./memory.loci",
        embedding_model: str = "all-MiniLM-L6-v2",
    ) -> None:
        self.path = Path(path)
        self._store = MemoryStore(self.path)
        self._indexer = Indexer(model_name=embedding_model)

    # ------------------------------------------------------------------
    # Write
    # ------------------------------------------------------------------

    def memorize(
        self,
        content: str,
        entities: Optional[List[str]] = None,
        scope: str = "global",
        importance: Optional[float] = None,
        relationships: Optional[List[Dict]] = None,
        type: str = "memory",
        resolve: Optional[str] = None,
    ) -> MemoryNode:
        """Persist a new memory.

        Parameters
        ----------
        content : str
            The raw text of the memory.
        entities : list[str], optional
            Entity tags to associate with this memory.
        scope : str
            Context scope (e.g. ``"project:dropzi"``).
        importance : float, optional
            Manual importance override.  If *None*, auto-scored.
        relationships : list[dict], optional
            List of ``{"target_entity": ..., "type": ...}`` dicts.
        type : str
            Memory type: ``"memory"``, ``"correction"``, ``"preference"``, or ``"fact"``.
        resolve : str, optional
            Node ID of an older memory this corrects. The old node is
            soft-deleted and a ``corrects`` edge is created.

        Returns
        -------
        MemoryNode
            The newly created node.
        """
        entities = entities or []
        embedding = self._indexer.embed(content)

        node = MemoryNode(
            content=content,
            embedding=embedding,
            entity_tags=entities,
            context_scope=scope,
            summary_level=0,
            memory_type=type,
        )

        if importance is not None:
            node.importance = importance
        else:
            node.importance = compute_importance(node, query_entities=entities)

        self._store.save_node(node)

        # If this corrects an older memory, soft-delete the old one and link
        if resolve:
            self._store.soft_delete_node(resolve)
            correction_edge = MemoryEdge(
                source_id=node.id,
                target_id=resolve,
                relationship_type="corrects",
            )
            self._store.save_edge(correction_edge)

        # Auto-create edges to nodes sharing entities
        if relationships:
            for rel in relationships:
                self._link_by_entity(
                    node, rel.get("target_entity", ""), rel.get("type", "related_to")
                )
        elif entities:
            self._auto_link(node, entities)

        return node

    def _auto_link(self, node: MemoryNode, entities: List[str]) -> None:
        """Create edges to the most recent nodes sharing any entity tag."""
        existing = self._store.list_nodes(scope=node.context_scope, limit=50)
        for other in existing:
            if other.id == node.id:
                continue
            overlap = set(node.entity_tags).intersection(other.entity_tags)
            if overlap:
                strength = len(overlap) / max(len(node.entity_tags), 1)
                edge = MemoryEdge(
                    source_id=node.id,
                    target_id=other.id,
                    relationship_type="related_to",
                    strength=min(strength, 1.0),
                )
                self._store.save_edge(edge)

    def _link_by_entity(
        self, node: MemoryNode, target_entity: str, rel_type: str
    ) -> None:
        """Create an edge from *node* to the most recent node tagged with *target_entity*."""
        candidates = self._store.list_nodes(scope=node.context_scope, limit=200)
        for c in candidates:
            if c.id == node.id:
                continue
            if target_entity in c.entity_tags:
                edge = MemoryEdge(
                    source_id=node.id,
                    target_id=c.id,
                    relationship_type=rel_type,
                )
                self._store.save_edge(edge)
                break  # link to most recent match only

    # ------------------------------------------------------------------
    # Read
    # ------------------------------------------------------------------

    def recall(
        self,
        query: str,
        entities: Optional[List[str]] = None,
        scope: Optional[str] = None,
        time_window_days: Optional[int] = None,
        min_importance: float = 0.0,
        top_k: int = 10,
        as_of: Optional[datetime] = None,
    ) -> List[Dict]:
        """Retrieve memories relevant to a query.

        Parameters
        ----------
        as_of : datetime, optional
            Time-travel: only return memories that existed at this point in time.

        Returns a list of dicts with keys ``node``, ``similarity``, and
        ``importance``.
        """
        query_embedding = self._indexer.embed(query)

        after = None
        if time_window_days is not None:
            after = datetime.now(timezone.utc) - timedelta(days=time_window_days)

        before = as_of  # time-travel cutoff

        results = self._store.vector_search(
            query_embedding=query_embedding,
            top_k=top_k * 3,  # over-fetch to re-rank
            scope=scope,
            min_importance=min_importance,
            entity_tags=entities,
            after=after,
            before=before,
        )

        # Re-rank using importance scorer
        scored: List[Tuple[MemoryNode, float, float]] = []
        for node, sim in results:
            imp = compute_importance(node, query_entities=entities)
            # Composite: 60% semantic similarity, 40% importance
            composite = 0.6 * sim + 0.4 * imp
            scored.append((node, sim, composite))

        scored.sort(key=lambda x: x[2], reverse=True)
        scored = scored[:top_k]

        # Touch accessed nodes
        for node, _, _ in scored:
            node.touch()
            self._store.update_node(node)

        return [
            {
                "node": node,
                "similarity": round(sim, 4),
                "importance": round(compute_importance(node, entities), 4),
                "composite_score": round(comp, 4),
            }
            for node, sim, comp in scored
        ]

    def get(self, node_id: str) -> Optional[MemoryNode]:
        """Retrieve a single memory node by ID."""
        return self._store.get_node(node_id)

    # ------------------------------------------------------------------
    # Graph
    # ------------------------------------------------------------------

    def traverse(
        self,
        from_entity: str,
        relationships: Optional[List[str]] = None,
        depth: int = 3,
    ) -> List[Dict]:
        """Walk the memory graph from nodes matching *from_entity*.

        Returns a list of ``{node, depth, path}`` dicts.
        """
        return self._store.traverse(
            start_entity=from_entity,
            relationship_types=relationships,
            max_depth=depth,
        )

    def link(
        self,
        source_id: str,
        target_id: str,
        relationship: str = "related_to",
        strength: float = 1.0,
    ) -> MemoryEdge:
        """Manually create an edge between two nodes."""
        edge = MemoryEdge(
            source_id=source_id,
            target_id=target_id,
            relationship_type=relationship,
            strength=strength,
        )
        self._store.save_edge(edge)
        return edge

    # ------------------------------------------------------------------
    # Maintenance
    # ------------------------------------------------------------------

    def compress(self, scope: Optional[str] = None, older_than_days: int = 2) -> int:
        """Run daily compression and return the count of new summary nodes."""
        summaries = compress_daily(self._store, scope=scope, older_than_days=older_than_days)
        return len(summaries)

    # ------------------------------------------------------------------
    # Export / Import
    # ------------------------------------------------------------------

    def export(
        self,
        path: str | Path,
        format: str = "native",
        max_tokens: int = 4000,
        scope: Optional[str] = None,
        min_importance: float = 0.0,
        time_window_days: Optional[int] = None,
    ) -> Path:
        """Export the database to a file.

        Parameters
        ----------
        path : str | Path
            Output file path.
        format : str
            ``"native"`` (full JSON), ``"context_prompt"`` (LLM-ready markdown),
            or ``"snapshot"`` (dense plain text).
        max_tokens : int
            Token budget for context_prompt and snapshot formats.
        scope : str, optional
            Filter to a specific scope (e.g. ``"project:dropzi"``).
            Pass ``"all"`` or *None* for all scopes.
        min_importance : float
            Minimum importance threshold.
        time_window_days : int, optional
            Only include memories from the last N days.

        Returns
        -------
        Path
            The path to the written file.
        """
        if scope == "all":
            scope = None

        if format == "native":
            return export_native(
                self._store, path, scope=scope,
                min_importance=min_importance,
                time_window_days=time_window_days,
            )
        elif format == "context_prompt":
            return export_context_prompt(
                self._store, path, max_tokens=max_tokens,
                scope=scope, min_importance=min_importance,
                time_window_days=time_window_days,
            )
        elif format == "snapshot":
            return export_snapshot(
                self._store, path, max_tokens=max_tokens,
                scope=scope, min_importance=min_importance,
                time_window_days=time_window_days,
            )
        else:
            raise ValueError(f"Unknown export format: {format!r}. Use 'native', 'context_prompt', or 'snapshot'.")

    def import_from(
        self,
        path: str | Path,
        mode: str = "replace",
    ) -> Dict:
        """Import memories from a file.

        Parameters
        ----------
        path : str | Path
            Path to a ``.loci.json`` or ``.context.md`` file.
        mode : str
            ``"replace"`` (clear DB first) or ``"merge"`` (add without overwriting).

        Returns
        -------
        dict
            Summary with imported_nodes count.
        """
        p = Path(path)
        suffix = p.suffix.lower()
        name_lower = p.name.lower()

        if name_lower.endswith(".loci.json") or suffix == ".json":
            return import_native(self._store, p, mode=mode)
        elif suffix == ".md":
            return import_context_md(self._store, p, scope="imported")
        else:
            # Try native JSON first
            try:
                return import_native(self._store, p, mode=mode)
            except (json.JSONDecodeError, KeyError):
                return import_context_md(self._store, p, scope="imported")

    def stats(self) -> Dict:
        """Return database statistics."""
        return self._store.stats()

    def forget(self, node_id: str, hard: bool = False) -> None:
        """Delete a memory node.

        By default performs a soft-delete (sets deleted_at timestamp).
        Pass ``hard=True`` to permanently remove the node and its edges.
        """
        if hard:
            self._store.delete_node(node_id)
        else:
            self._store.soft_delete_node(node_id)

    # ------------------------------------------------------------------
    # Persona
    # ------------------------------------------------------------------

    def set_persona(self, key: str, value: str) -> None:
        """Set a single persona attribute (e.g. name, role, communication_style)."""
        _set_persona(self._store, key, value)

    def get_persona(self, key: Optional[str] = None) -> Any:
        """Get a persona attribute by key, or all attributes if key is None."""
        if key is None:
            return _get_all_persona(self._store)
        return _get_persona(self._store, key)

    def update_persona(self, updates: Dict[str, str]) -> None:
        """Batch-update persona attributes."""
        _update_persona(self._store, updates)

    def clear_persona(self) -> None:
        """Remove all persona data."""
        _clear_persona(self._store)

    def get_persona_prompt(self) -> str:
        """Format persona data as a system prompt prefix."""
        return _format_persona_prompt(self._store)

    # ------------------------------------------------------------------
    # Corrections
    # ------------------------------------------------------------------

    def correct(
        self,
        old_node_id: str,
        new_content: str,
        entities: Optional[List[str]] = None,
        scope: str = "global",
    ) -> MemoryNode:
        """Create a correction that supersedes an older memory.

        The old node is soft-deleted, and a ``corrects`` edge is created.
        """
        return self.memorize(
            content=new_content,
            entities=entities,
            scope=scope,
            type="correction",
            resolve=old_node_id,
        )

    def get_corrections(self, scope: Optional[str] = None, limit: int = 50) -> List[MemoryNode]:
        """List all correction-type memories."""
        return self._store.list_nodes(
            scope=scope,
            min_importance=0.0,
            limit=limit,
            memory_type="correction",
        )

    # ------------------------------------------------------------------
    # Audit Trail
    # ------------------------------------------------------------------

    def store_trace(
        self,
        query: str,
        node_ids_used: List[str],
        response: str = "",
        scope: str = "global",
        metadata: Optional[Dict] = None,
    ) -> str:
        """Record a reasoning trace. Returns the trace ID."""
        return _store_trace(
            self._store,
            query=query,
            node_ids_used=node_ids_used,
            response=response,
            scope=scope,
            metadata=metadata,
        )

    def explain(self, trace_id: str) -> Optional[str]:
        """Get a human-readable explanation of a trace."""
        return _explain_trace(self._store, trace_id)

    def traces(self, scope: Optional[str] = None, limit: int = 50) -> List[Dict]:
        """List recent reasoning traces."""
        return _list_traces(self._store, scope=scope, limit=limit)

    def traces_using_node(self, node_id: str) -> List[Dict]:
        """Find all traces that used a specific memory node."""
        return _traces_using_node(self._store, node_id)

    # ------------------------------------------------------------------
    # Contradiction Detection
    # ------------------------------------------------------------------

    def check_conflicts(
        self,
        scope: Optional[str] = None,
        similarity_threshold: float = 0.85,
    ) -> List[ConflictWarning]:
        """Scan all memories for contradictions."""
        return _check_conflicts(
            self._store,
            scope=scope,
            similarity_threshold=similarity_threshold,
        )

    # ------------------------------------------------------------------
    # Time-Travel
    # ------------------------------------------------------------------

    def snapshot(self, as_of: datetime, scope: Optional[str] = None) -> List[MemoryNode]:
        """Return all memories as they existed at a point in time.

        Includes nodes created before ``as_of`` that were not yet deleted,
        or were deleted after ``as_of``.
        """
        all_nodes = self._store.list_nodes(
            scope=scope,
            min_importance=0.0,
            limit=10000,
            include_deleted=True,
        )
        result = []
        for n in all_nodes:
            if n.timestamp > as_of:
                continue  # created after the snapshot point
            if n.deleted_at and n.deleted_at <= as_of:
                continue  # already deleted at that point
            result.append(n)
        return result

    def timeline(
        self,
        entity: Optional[str] = None,
        scope: Optional[str] = None,
        limit: int = 100,
    ) -> List[MemoryNode]:
        """Return memories in chronological order, optionally filtered by entity."""
        nodes = self._store.list_nodes(
            scope=scope,
            min_importance=0.0,
            limit=limit,
            include_deleted=True,
        )
        if entity:
            nodes = [n for n in nodes if entity in n.entity_tags]
        nodes.sort(key=lambda n: n.timestamp)
        return nodes

    def diff(
        self,
        t1: datetime,
        t2: datetime,
        scope: Optional[str] = None,
    ) -> Dict:
        """Compare memory state between two points in time.

        Returns dict with ``added``, ``removed``, and ``modified`` lists.
        """
        snap1 = {n.id: n for n in self.snapshot(t1, scope=scope)}
        snap2 = {n.id: n for n in self.snapshot(t2, scope=scope)}

        added = [snap2[nid] for nid in snap2 if nid not in snap1]
        removed = [snap1[nid] for nid in snap1 if nid not in snap2]

        return {
            "added": added,
            "removed": removed,
            "t1_count": len(snap1),
            "t2_count": len(snap2),
        }

    # ------------------------------------------------------------------
    # Health
    # ------------------------------------------------------------------

    def health(self, scope: Optional[str] = None) -> HealthReport:
        """Compute a health score for the memory database."""
        return compute_health(self._store, scope=scope)

    def get_onboarding_questions(self) -> List[str]:
        """Generate questions to help bootstrap the memory database."""
        return _get_onboarding_questions(self._store)

    def close(self) -> None:
        """Close the underlying store."""
        self._store.close()

    # Context manager support
    def __enter__(self) -> "LociDB":
        return self

    def __exit__(self, *exc) -> None:
        self.close()
