"""Memory portability — export and import for LociDB.

Supports three export formats:
  - **native**:         Full-fidelity JSON (.loci.json) for backup/migration.
  - **context_prompt**: Structured markdown (.context.md) for LLM system prompts.
  - **snapshot**:       Dense plain-text bullet points (.snapshot.txt).

And import with replace or merge semantics.
"""

from __future__ import annotations

import json
import re
from collections import defaultdict
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Dict, List, Optional

from .edge import MemoryEdge
from .node import MemoryNode

if TYPE_CHECKING:
    from .store import MemoryStore

__all__ = ["export_native", "export_context_prompt", "export_snapshot", "import_native", "import_context_md"]

# ── Helpers ───────────────────────────────────────────────────────────

_ROUGH_CHARS_PER_TOKEN = 4  # conservative estimate for English text


def _estimate_tokens(text: str) -> int:
    return len(text) // _ROUGH_CHARS_PER_TOKEN


def _fetch_nodes(
    store: "MemoryStore",
    scope: Optional[str],
    min_importance: float,
    time_window_days: Optional[int],
) -> List[MemoryNode]:
    """Fetch filtered, importance-sorted nodes from the store."""
    nodes = store.list_nodes(scope=scope, min_importance=min_importance, limit=100_000)
    if time_window_days is not None:
        cutoff = datetime.now(timezone.utc) - timedelta(days=time_window_days)
        nodes = [
            n for n in nodes
            if (n.timestamp.replace(tzinfo=timezone.utc) if n.timestamp.tzinfo is None else n.timestamp) >= cutoff
        ]
    # Sort by importance descending (deterministic: break ties by id)
    nodes.sort(key=lambda n: (-n.importance, n.id))
    return nodes


def _fetch_all_edges(store: "MemoryStore") -> List[MemoryEdge]:
    """Fetch every edge in the store (used by native export)."""
    cur = store._conn.cursor()
    cur.execute("SELECT * FROM edges")
    return [store._row_to_edge(r) for r in cur.fetchall()]


def _fetch_edges_for_nodes(store: "MemoryStore", node_ids: set) -> List[MemoryEdge]:
    """Fetch edges where both endpoints are in *node_ids*."""
    all_edges = _fetch_all_edges(store)
    return [e for e in all_edges if e.source_id in node_ids and e.target_id in node_ids]


# ── Native export / import ────────────────────────────────────────────

def export_native(
    store: "MemoryStore",
    path: str | Path,
    scope: Optional[str] = None,
    min_importance: float = 0.0,
    time_window_days: Optional[int] = None,
) -> Path:
    """Export full-fidelity JSON."""
    out = Path(path)
    nodes = _fetch_nodes(store, scope, min_importance, time_window_days)
    node_ids = {n.id for n in nodes}
    edges = _fetch_edges_for_nodes(store, node_ids)

    scopes = sorted({n.context_scope for n in nodes})

    payload = {
        "locidb_version": "0.1.0",
        "exported_at": datetime.now(timezone.utc).isoformat(),
        "stats": {
            "total_nodes": len(nodes),
            "total_edges": len(edges),
            "scopes": scopes,
        },
        "nodes": [
            {
                "id": n.id,
                "content": n.content,
                "embedding": n.embedding,
                "importance": n.importance,
                "timestamp": n.timestamp.isoformat(),
                "entity_tags": n.entity_tags,
                "context_scope": n.context_scope,
                "summary_level": n.summary_level,
                "access_count": n.access_count,
                "last_accessed": n.last_accessed.isoformat() if n.last_accessed else None,
            }
            for n in nodes
        ],
        "edges": [
            {
                "id": e.id,
                "source_id": e.source_id,
                "target_id": e.target_id,
                "relationship_type": e.relationship_type,
                "strength": e.strength,
                "timestamp": e.timestamp.isoformat(),
            }
            for e in edges
        ],
    }

    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    return out


def import_native(
    store: "MemoryStore",
    path: str | Path,
    mode: str = "replace",
) -> Dict:
    """Import a native .loci.json file.

    *mode*:
      - ``"replace"``: clear existing data first.
      - ``"merge"``:   add new nodes/edges, skip existing IDs.
    """
    data = json.loads(Path(path).read_text(encoding="utf-8"))

    if mode == "replace":
        cur = store._conn.cursor()
        cur.execute("DELETE FROM edges")
        cur.execute("DELETE FROM nodes")
        store._conn.commit()

    imported_nodes = 0
    imported_edges = 0

    for nd in data.get("nodes", []):
        if mode == "merge" and store.get_node(nd["id"]) is not None:
            continue
        node = MemoryNode(
            id=nd["id"],
            content=nd["content"],
            embedding=nd.get("embedding", []),
            importance=nd.get("importance", 0.5),
            timestamp=datetime.fromisoformat(nd["timestamp"]),
            entity_tags=nd.get("entity_tags", []),
            context_scope=nd.get("context_scope", "global"),
            summary_level=nd.get("summary_level", 0),
            access_count=nd.get("access_count", 0),
            last_accessed=(
                datetime.fromisoformat(nd["last_accessed"])
                if nd.get("last_accessed")
                else None
            ),
        )
        store.save_node(node)
        imported_nodes += 1

    for ed in data.get("edges", []):
        if mode == "merge":
            # Check if edge already exists
            cur = store._conn.cursor()
            cur.execute("SELECT 1 FROM edges WHERE id = ?", (ed["id"],))
            if cur.fetchone() is not None:
                continue
        edge = MemoryEdge(
            id=ed["id"],
            source_id=ed["source_id"],
            target_id=ed["target_id"],
            relationship_type=ed["relationship_type"],
            strength=ed.get("strength", 1.0),
            timestamp=datetime.fromisoformat(ed["timestamp"]) if "timestamp" in ed else datetime.now(timezone.utc),
        )
        store.save_edge(edge)
        imported_edges += 1

    return {"imported_nodes": imported_nodes, "imported_edges": imported_edges, "mode": mode}


# ── Context prompt export ─────────────────────────────────────────────

def export_context_prompt(
    store: "MemoryStore",
    path: str | Path,
    max_tokens: int = 4000,
    scope: Optional[str] = None,
    min_importance: float = 0.0,
    time_window_days: Optional[int] = None,
) -> Path:
    """Export a structured markdown file suitable as an LLM system prompt."""
    out = Path(path)
    nodes = _fetch_nodes(store, scope, min_importance, time_window_days)
    node_ids = {n.id for n in nodes}
    edges = _fetch_edges_for_nodes(store, node_ids)

    now = datetime.now(timezone.utc)
    date_str = now.strftime("%Y-%m-%d %H:%M UTC")

    # Prefer compressed summaries over raw events
    nodes.sort(key=lambda n: (-n.summary_level, -n.importance, n.id))

    # Budget-aware selection
    selected: List[MemoryNode] = []
    token_budget = max_tokens - 300  # reserve for headers/footer
    used = 0
    for n in nodes:
        est = _estimate_tokens(n.content) + 10  # overhead per entry
        if used + est > token_budget:
            break
        selected.append(n)
        used += est

    # Group by scope then by entity
    by_scope: Dict[str, List[MemoryNode]] = defaultdict(list)
    for n in selected:
        by_scope[n.context_scope].append(n)

    # Identify "identity" nodes (user:* scopes)
    identity_nodes = []
    project_scopes: Dict[str, List[MemoryNode]] = {}
    other_nodes: List[MemoryNode] = []
    for scope_key, scope_nodes in sorted(by_scope.items()):
        if scope_key.startswith("user:"):
            identity_nodes.extend(scope_nodes)
        elif scope_key.startswith("project:"):
            project_scopes[scope_key] = scope_nodes
        else:
            other_nodes.extend(scope_nodes)

    # Recent = last 30 days
    cutoff_30 = now - timedelta(days=30)
    recent = [
        n for n in selected
        if (n.timestamp.replace(tzinfo=timezone.utc) if n.timestamp.tzinfo is None else n.timestamp) >= cutoff_30
    ]

    # Key decisions = high importance, low summary_level
    decisions = [n for n in selected if n.importance >= 0.7]
    decisions.sort(key=lambda n: -n.importance)

    # Build relationship highlights
    rel_lines: List[str] = []
    node_map = {n.id: n for n in selected}
    seen_rels: set = set()
    for e in edges:
        if e.source_id in node_map and e.target_id in node_map:
            key = (e.source_id, e.target_id, e.relationship_type)
            if key in seen_rels:
                continue
            seen_rels.add(key)
            src = node_map[e.source_id]
            tgt = node_map[e.target_id]
            src_label = src.entity_tags[0] if src.entity_tags else src.content[:30]
            tgt_label = tgt.entity_tags[0] if tgt.entity_tags else tgt.content[:30]
            rel_lines.append(f"- **{src_label}** —[{e.relationship_type}]→ **{tgt_label}**")

    # Assemble markdown
    lines: List[str] = []
    lines.append(f"# Memory Context")
    lines.append(f"*Exported from LociDB on {date_str} | {len(selected)} memories distilled*")
    lines.append("")

    if identity_nodes:
        lines.append("## Core identity & background")
        for n in identity_nodes:
            lines.append(f"- {n.content}")
        lines.append("")

    if project_scopes:
        lines.append("## Active projects")
        for scope_key, scope_nodes in sorted(project_scopes.items()):
            project_name = scope_key.split(":", 1)[1] if ":" in scope_key else scope_key
            lines.append(f"### {project_name}")
            for n in scope_nodes:
                lines.append(f"- {n.content}")
            lines.append("")

    if decisions:
        lines.append("## Key decisions & preferences")
        for n in decisions[:15]:
            imp_str = f"[{n.importance:.2f}]"
            lines.append(f"- {imp_str} {n.content}")
        lines.append("")

    if recent:
        lines.append("## Recent context (last 30 days)")
        for n in recent[:15]:
            ts_str = n.timestamp.strftime("%Y-%m-%d")
            lines.append(f"- [{ts_str}] {n.content}")
        lines.append("")

    if other_nodes:
        lines.append("## General memories")
        for n in other_nodes:
            lines.append(f"- {n.content}")
        lines.append("")

    if rel_lines:
        lines.append("## Knowledge graph highlights")
        lines.extend(rel_lines[:20])
        lines.append("")

    lines.append("---")
    lines.append("*To use: paste this as a system prompt or project instruction in any LLM*")
    lines.append("*To restore full memory: import the .loci.json file into LociDB*")

    content = "\n".join(lines) + "\n"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(content, encoding="utf-8")
    return out


# ── Snapshot export ───────────────────────────────────────────────────

def export_snapshot(
    store: "MemoryStore",
    path: str | Path,
    max_tokens: int = 2000,
    scope: Optional[str] = None,
    min_importance: float = 0.0,
    time_window_days: Optional[int] = None,
) -> Path:
    """Export a dense plain-text snapshot of top memories."""
    out = Path(path)
    nodes = _fetch_nodes(store, scope, min_importance, time_window_days)
    # Prefer summaries
    nodes.sort(key=lambda n: (-n.summary_level, -n.importance, n.id))

    lines: List[str] = []
    used = 0
    budget = max_tokens - 50  # small header
    for n in nodes:
        bullet = f"• {n.content}"
        est = _estimate_tokens(bullet) + 2
        if used + est > budget:
            break
        lines.append(bullet)
        used += est

    header = f"LociDB snapshot — {len(lines)} memories\n"
    content = header + "\n".join(lines) + "\n"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(content, encoding="utf-8")
    return out


# ── Context markdown import ──────────────────────────────────────────

def import_context_md(
    store: "MemoryStore",
    path: str | Path,
    scope: str = "imported",
) -> Dict:
    """Parse a markdown context file and store each bullet as a memory.

    This is a best-effort import — it extracts bullet lines and stores
    them as raw (summary_level=0) memories.  For full fidelity, use
    the native .loci.json format.
    """
    text = Path(path).read_text(encoding="utf-8")

    # Extract lines starting with "- " or "• "
    bullets = re.findall(r"^[\-•]\s+(.+)$", text, re.MULTILINE)

    imported = 0
    for content in bullets:
        # Strip optional leading tags like [0.87] or [2026-03-20]
        clean = re.sub(r"^\[[\d.\-:\sA-Za-z]+\]\s*", "", content).strip()
        if not clean:
            continue
        node = MemoryNode(
            content=clean,
            context_scope=scope,
            summary_level=0,
        )
        store.save_node(node)
        imported += 1

    return {"imported_nodes": imported, "source": str(path)}
