"""
cortex/tools/migrate_mem0.py — Mem0 API → Cortex SQLite migration tool.

Imports memories from the Mem0 Platform API into a Cortex SQLite database.
Designed to be idempotent (safe to re-run), with optional dry-run mode.

Usage::

    python -m cortex.tools.migrate_mem0 \\
        --mem0-api-key  $MEM0_API_KEY \\
        --mem0-org-id   $MEM0_ORG_ID \\
        --mem0-project-id $MEM0_PROJECT_ID \\
        --agent-id      eva \\
        --cortex-db     ./cortex.db \\
        --dry-run

Field Mapping (Mem0 → Cortex semantic_claim):
    memory           → object_value   (the memory text)
    id               → source_session (as "mem0:<id>" for idempotency)
    metadata.kind    → predicate      (fallback: "recalls")
    metadata.subject → subject        (fallback: agent_id or user_id)
    created_at       → created_at     (ISO 8601 UTC)

Cornerstones:
    Memories with metadata.is_cornerstone == "true" are additionally sealed
    via CornerstoneGuardian.seal() with a golden hash for drift detection.

Idempotency:
    Each imported claim is stored with source_session = "mem0:<mem0_id>".
    On re-run, the tool queries for that value and skips already-imported IDs.

Dependencies:
    httpx (API calls) — already in pyproject.toml
    aiosqlite (via SQLiteAdapter)
    cortex.storage.sqlite_adapter.SQLiteAdapter
    cortex.identity.cornerstone.CornerstoneGuardian
    cortex.config.CortexConfig
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import re
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Tuple

try:
    import httpx
except ImportError:
    sys.exit("httpx is required: pip install httpx")

from cortex.config import CortexConfig
from cortex.identity.cornerstone import CornerstoneGuardian
from cortex.storage.sqlite_adapter import SQLiteAdapter

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

MEM0_API_BASE = "https://api.mem0.ai/v1"
MEM0_SOURCE_PREFIX = "mem0:"  # stored in source_session for idempotency
DEFAULT_PREDICATE = "recalls"
DEFAULT_SUBJECT_FALLBACK = "agent"

# ---------------------------------------------------------------------------
# Mem0 memory representation
# ---------------------------------------------------------------------------


@dataclass
class Mem0Memory:
    """Lightweight representation of a single Mem0 API memory record."""

    id: str
    memory: str
    agent_id: Optional[str] = None
    user_id: Optional[str] = None
    created_at: Optional[str] = None
    updated_at: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_api_dict(cls, data: Dict[str, Any]) -> "Mem0Memory":
        """Parse a raw Mem0 API response dict into a Mem0Memory."""
        return cls(
            id=data["id"],
            memory=data.get("memory", ""),
            agent_id=data.get("agent_id"),
            user_id=data.get("user_id"),
            created_at=data.get("created_at"),
            updated_at=data.get("updated_at"),
            metadata=data.get("metadata") or {},
        )

    @property
    def is_cornerstone(self) -> bool:
        """True if this memory should be sealed as a Cortex cornerstone."""
        return str(self.metadata.get("is_cornerstone", "false")).lower() == "true"

    @property
    def source_session_key(self) -> str:
        """Unique idempotency key stored in semantic_claim.source_session."""
        return f"{MEM0_SOURCE_PREFIX}{self.id}"


# ---------------------------------------------------------------------------
# Mem0 API client
# ---------------------------------------------------------------------------


class Mem0Client:
    """Minimal async Mem0 API client using httpx."""

    def __init__(
        self,
        api_key: str,
        org_id: Optional[str] = None,
        project_id: Optional[str] = None,
    ) -> None:
        self._api_key = api_key
        self._org_id = org_id
        self._project_id = project_id
        self._headers: Dict[str, str] = {
            "Authorization": f"Token {api_key}",
            "Content-Type": "application/json",
        }
        if org_id:
            self._headers["Mem0-Org-Id"] = org_id
        if project_id:
            self._headers["Mem0-Project-Id"] = project_id

    async def fetch_memories(
        self,
        agent_id: Optional[str] = None,
        user_id: Optional[str] = None,
        page_size: int = 100,
    ) -> List[Mem0Memory]:
        """
        Fetch all memories for the given agent_id or user_id.

        Paginates automatically until all pages are retrieved.

        Raises:
            ValueError: if neither agent_id nor user_id is provided.
            httpx.HTTPStatusError: on non-2xx responses.
        """
        if not agent_id and not user_id:
            raise ValueError("Either agent_id or user_id must be provided")

        params: Dict[str, Any] = {"page_size": page_size, "page": 1}
        if agent_id:
            params["agent_id"] = agent_id
        if user_id:
            params["user_id"] = user_id

        memories: List[Mem0Memory] = []

        async with httpx.AsyncClient(timeout=30.0) as client:
            while True:
                resp = await client.get(
                    f"{MEM0_API_BASE}/memories/",
                    headers=self._headers,
                    params=params,
                )
                resp.raise_for_status()
                payload = resp.json()

                # Mem0 API may return a list directly or a paginated envelope
                if isinstance(payload, list):
                    batch = [Mem0Memory.from_api_dict(m) for m in payload]
                    memories.extend(batch)
                    break  # list response = single page
                else:
                    results = payload.get("results", payload.get("memories", []))
                    batch = [Mem0Memory.from_api_dict(m) for m in results]
                    memories.extend(batch)

                    # Check for next page
                    if not payload.get("next") or len(batch) < page_size:
                        break
                    params["page"] = params["page"] + 1

        return memories


# ---------------------------------------------------------------------------
# Field mapping helpers
# ---------------------------------------------------------------------------


def _map_subject(mem: Mem0Memory, fallback: str) -> str:
    """Derive claim subject from Mem0 metadata or identity fields."""
    if mem.metadata.get("subject"):
        return str(mem.metadata["subject"])
    if mem.agent_id:
        return mem.agent_id
    if mem.user_id:
        return mem.user_id
    return fallback


def _map_predicate(mem: Mem0Memory) -> str:
    """Derive claim predicate from Mem0 metadata.kind or default."""
    kind = mem.metadata.get("kind") or mem.metadata.get("type")
    if kind:
        return str(kind)
    return DEFAULT_PREDICATE


def _map_created_at(mem: Mem0Memory) -> Optional[str]:
    """Return Mem0's created_at if present, else None (SQLiteAdapter will use _now())."""
    ts = mem.created_at
    if not ts:
        return None
    # Normalise to UTC ISO 8601 with milliseconds if needed
    try:
        dt = datetime.fromisoformat(ts.replace("Z", "+00:00"))
        return dt.astimezone(timezone.utc).isoformat(timespec="milliseconds")
    except ValueError:
        return ts  # pass through as-is if unparseable


def _extract_entities(text: str) -> List[Tuple[str, str]]:
    """
    Lightweight named-entity heuristic: extract capitalised multi-word sequences.

    Returns a list of (canonical_name, entity_type) tuples.
    entity_type is 'person' or 'project' (heuristic, not guaranteed).

    This is intentionally simple — production use should wire an LLMProvider
    for proper NER. The heuristic avoids false positives from sentence-start
    capitalisation by requiring at least two consecutive title-cased tokens.
    """
    # Tokenise and look for consecutive Title-Case words (likely proper nouns)
    tokens = re.findall(r"\b([A-Z][a-z]{1,}(?:\s+[A-Z][a-z]{1,})*)\b", text)
    seen: set[str] = set()
    entities: List[Tuple[str, str]] = []
    for tok in tokens:
        if tok in seen or len(tok) < 4:  # skip very short tokens / duplicates
            continue
        seen.add(tok)
        # Crude type heuristic: if it looks like a project/tool name use "project"
        etype = "project" if any(
            kw in tok.lower() for kw in ("api", "sdk", "lib", "tool", "app", "bot", "cli")
        ) else "person"
        entities.append((tok, etype))
    return entities


# ---------------------------------------------------------------------------
# Migration result summary
# ---------------------------------------------------------------------------


@dataclass
class MigrationReport:
    """Summary of a completed (or dry-run) migration."""

    total_fetched: int = 0
    imported_claims: int = 0
    skipped_duplicates: int = 0
    imported_cornerstones: int = 0
    imported_entities: int = 0
    errors: List[str] = field(default_factory=list)
    dry_run: bool = False

    def __str__(self) -> str:
        prefix = "[DRY-RUN] " if self.dry_run else ""
        return (
            f"{prefix}Migration complete.\n"
            f"  Fetched from Mem0 : {self.total_fetched}\n"
            f"  Claims imported   : {self.imported_claims}\n"
            f"  Cornerstones      : {self.imported_cornerstones}\n"
            f"  Entities          : {self.imported_entities}\n"
            f"  Skipped (dup)     : {self.skipped_duplicates}\n"
            f"  Errors            : {len(self.errors)}"
        )


# ---------------------------------------------------------------------------
# Core migrator
# ---------------------------------------------------------------------------


class Mem0Migrator:
    """
    Orchestrates the Mem0 → Cortex migration.

    Typical usage::

        async with SQLiteAdapter(db_path) as adapter:
            migrator = Mem0Migrator(adapter, owner_id="default", dry_run=False)
            report = await migrator.run(memories)

    The migrator is stateless between runs — idempotency is implemented via
    the ``source_session`` column (value = "mem0:<mem0_id>").
    """

    def __init__(
        self,
        adapter: SQLiteAdapter,
        owner_id: str = "default",
        dry_run: bool = False,
        subject_fallback: str = DEFAULT_SUBJECT_FALLBACK,
        create_entities: bool = True,
    ) -> None:
        self._adapter = adapter
        self._owner_id = owner_id
        self._dry_run = dry_run
        self._subject_fallback = subject_fallback
        self._create_entities = create_entities
        # Lazy-initialised cornerstone guardian (needs config)
        self._guardian: Optional[CornerstoneGuardian] = None

    def _get_guardian(self) -> CornerstoneGuardian:
        if self._guardian is None:
            config = CortexConfig()
            self._guardian = CornerstoneGuardian(self._adapter, config)
        return self._guardian

    # ------------------------------------------------------------------ #
    # Idempotency check
    # ------------------------------------------------------------------ #

    async def _is_already_imported(self, mem: Mem0Memory) -> bool:
        """
        Return True if this Mem0 memory was already imported.

        We use source_session = "mem0:<id>" as the idempotency key and look
        for any claim with that source_session in the owner's namespace.
        """
        try:
            # get_session_claims queries WHERE source_session = ?
            existing = await self._adapter.get_session_claims(mem.source_session_key)
            # Filter to this owner
            return any(c.owner_id == self._owner_id for c in existing)
        except Exception:
            return False

    # ------------------------------------------------------------------ #
    # Claim import
    # ------------------------------------------------------------------ #

    async def _import_claim(self, mem: Mem0Memory) -> Optional[str]:
        """
        Import a single Mem0 memory as a Cortex semantic_claim.

        Returns the new claim ID, or None if skipped (dry-run).
        """
        subject = _map_subject(mem, self._subject_fallback)
        predicate = _map_predicate(mem)
        object_value = mem.memory
        created_at = _map_created_at(mem)

        logger.debug(
            "Importing claim: subject=%r predicate=%r mem0_id=%s",
            subject, predicate, mem.id,
        )

        if self._dry_run:
            return None

        claim = await self._adapter.create_claim(
            subject=subject,
            predicate=predicate,
            object_value=object_value,
            owner_id=self._owner_id,
            memory_type="semantic",
            confidence=float(mem.metadata.get("confidence", 0.8)),
            source_session=mem.source_session_key,
            # Preserve original Mem0 created_at if available
            **({"created_at": created_at} if created_at else {}),
        )
        return claim.id

    # ------------------------------------------------------------------ #
    # Cornerstone import
    # ------------------------------------------------------------------ #

    async def _import_cornerstone(self, mem: Mem0Memory, claim_id: Optional[str]) -> bool:
        """
        Seal a Mem0 cornerstone memory via CornerstoneGuardian.

        Returns True if sealed (or would be in non-dry-run), False on error.
        """
        label = (
            mem.metadata.get("label")
            or mem.metadata.get("kind")
            or f"mem0:{mem.id[:8]}"
        )
        content = mem.memory

        logger.debug("Sealing cornerstone: label=%r mem0_id=%s", label, mem.id)

        if self._dry_run:
            return True

        try:
            guardian = self._get_guardian()
            await guardian.seal(
                label=str(label),
                content=content,
                owner_id=self._owner_id,
                sealed_by=f"migrate_mem0:{mem.id}",
            )
            return True
        except Exception as exc:
            logger.warning("Failed to seal cornerstone mem0_id=%s: %s", mem.id, exc)
            return False

    # ------------------------------------------------------------------ #
    # Entity extraction + import
    # ------------------------------------------------------------------ #

    async def _import_entities(self, mem: Mem0Memory) -> int:
        """
        Extract named entities from a Mem0 memory and store them.

        Returns the count of new entities created.
        """
        if not self._create_entities:
            return 0

        entities = _extract_entities(mem.memory)
        count = 0

        for canonical_name, entity_type in entities:
            logger.debug("Entity candidate: %r (%s)", canonical_name, entity_type)

            if self._dry_run:
                count += 1
                continue

            # Avoid duplicates: check alias table first
            try:
                existing = await self._adapter.find_entity_by_alias(
                    canonical_name, self._owner_id
                )
                if existing:
                    continue  # already known

                entity = await self._adapter.create_entity(
                    canonical_name=canonical_name,
                    entity_type=entity_type,
                    owner_id=self._owner_id,
                )
                await self._adapter.add_entity_alias(
                    entity_id=entity.id,
                    alias=canonical_name,
                    alias_type="name",
                    source="migrate_mem0",
                )
                count += 1
            except Exception as exc:
                logger.debug("Entity import failed for %r: %s", canonical_name, exc)

        return count

    # ------------------------------------------------------------------ #
    # Main run loop
    # ------------------------------------------------------------------ #

    async def run(self, memories: List[Mem0Memory]) -> MigrationReport:
        """
        Migrate a list of Mem0Memory objects into Cortex SQLite.

        Args:
            memories: Pre-fetched list of Mem0 memories to import.

        Returns:
            MigrationReport with counts and any errors.
        """
        report = MigrationReport(
            total_fetched=len(memories),
            dry_run=self._dry_run,
        )

        for mem in memories:
            try:
                # --- Idempotency check ---
                if not self._dry_run and await self._is_already_imported(mem):
                    logger.info("Skipping duplicate mem0_id=%s", mem.id)
                    report.skipped_duplicates += 1
                    continue

                # --- Import claim ---
                claim_id = await self._import_claim(mem)
                report.imported_claims += 1

                # --- Cornerstone ---
                if mem.is_cornerstone:
                    sealed = await self._import_cornerstone(mem, claim_id)
                    if sealed:
                        report.imported_cornerstones += 1
                    else:
                        report.errors.append(
                            f"Cornerstone seal failed for mem0_id={mem.id}"
                        )

                # --- Entities ---
                n_entities = await self._import_entities(mem)
                report.imported_entities += n_entities

            except Exception as exc:
                msg = f"Error importing mem0_id={mem.id}: {exc}"
                logger.error(msg)
                report.errors.append(msg)

        return report


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="python -m cortex.tools.migrate_mem0",
        description="Migrate memories from Mem0 API into Cortex SQLite.",
    )
    parser.add_argument(
        "--mem0-api-key",
        required=True,
        help="Mem0 API key (or set MEM0_API_KEY env var)",
    )
    parser.add_argument("--mem0-org-id", help="Mem0 organisation ID")
    parser.add_argument("--mem0-project-id", help="Mem0 project ID")

    id_group = parser.add_mutually_exclusive_group(required=True)
    id_group.add_argument("--agent-id", help="Mem0 agent_id to migrate")
    id_group.add_argument("--user-id", help="Mem0 user_id to migrate")

    parser.add_argument(
        "--cortex-db",
        default="./cortex.db",
        help="Path to Cortex SQLite database (created if absent). Default: ./cortex.db",
    )
    parser.add_argument(
        "--owner-id",
        default="default",
        help="Cortex owner_id namespace for imported memories. Default: default",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Preview what would be imported without writing anything.",
    )
    parser.add_argument(
        "--no-entities",
        action="store_true",
        help="Skip entity extraction and import.",
    )
    parser.add_argument(
        "--page-size",
        type=int,
        default=100,
        help="Mem0 API page size. Default: 100",
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Enable debug logging.",
    )
    return parser


async def _async_main(args: argparse.Namespace) -> int:
    """Async main — returns exit code."""
    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO, format="%(message)s")

    # ---- Fetch from Mem0 ----
    print(f"Connecting to Mem0 API…")
    client = Mem0Client(
        api_key=args.mem0_api_key,
        org_id=args.mem0_org_id,
        project_id=args.mem0_project_id,
    )

    try:
        memories = await client.fetch_memories(
            agent_id=getattr(args, "agent_id", None),
            user_id=getattr(args, "user_id", None),
            page_size=args.page_size,
        )
    except httpx.HTTPStatusError as exc:
        print(f"Error: Mem0 API returned HTTP {exc.response.status_code}: {exc.response.text}")
        return 1
    except Exception as exc:
        print(f"Error fetching from Mem0: {exc}")
        return 1

    print(f"Fetched {len(memories)} memories from Mem0.")

    if not memories:
        print("Nothing to migrate.")
        return 0

    if args.dry_run:
        print("[DRY-RUN] No data will be written to Cortex.")

    # ---- Open Cortex DB ----
    async with SQLiteAdapter(db_path=args.cortex_db) as adapter:
        migrator = Mem0Migrator(
            adapter=adapter,
            owner_id=args.owner_id,
            dry_run=args.dry_run,
            subject_fallback=getattr(args, "agent_id", None) or getattr(args, "user_id", "agent"),
            create_entities=not args.no_entities,
        )
        report = await migrator.run(memories)

    print(str(report))

    if report.errors:
        print("\nErrors:")
        for err in report.errors:
            print(f"  - {err}")
        return 1

    return 0


def main() -> None:
    parser = _build_parser()
    args = parser.parse_args()
    sys.exit(asyncio.run(_async_main(args)))


if __name__ == "__main__":
    main()
