"""
export.py
---------
Export tools for training data generation.
"""
from typing import Any
from datetime import datetime
import json
from hieromem_core.memory import Memory


def export_training_data(
    memory: Memory,
    namespace: str | None = None,
    format: str = "jsonl",
    include_graph: bool = True,
    start: datetime | None = None,
    end: datetime | None = None,
) -> Any:
    """
    Export memory episodes for fine-tuning or analysis.

    Args:
        memory: Memory instance
        namespace: Namespace/Realm to export from
        format: 'jsonl' (return string) or 'list' (return list of dicts)
        include_graph: Whether to include graph structure in export
    """
    # Reuse replay generator for filtering/sorting
    episodes = memory.replay(start=start, end=end, namespace=namespace)

    data = []
    for ep in episodes:
        record = {
            "id": ep.id,
            "timestamp": ep.tocc.isoformat(),
            "embedding": ep.e.tolist() if ep.e is not None else [],
            "content": ep.meta.get("text", "") if ep.meta else "",
            "meta": ep.meta,
        }

        if include_graph and ep.G:
            record["graph"] = {
                "nodes": [{"id": n, **(ep.G.nodes[n] or {})} for n in ep.G.nodes],
                "edges": [{"u": u, "v": v, **(ep.G.edges[u, v] or {})} for u, v in ep.G.edges],
            }

        data.append(record)

    if format == "jsonl":
        return "\n".join(json.dumps(d) for d in data)
    elif format == "list":
        return data
    else:
        raise ValueError(f"Unknown format: {format}")
