# SPDX-FileCopyrightText: 2026 Surendra Singh <66422685+bolnet@users.noreply.github.com>
# SPDX-License-Identifier: MIT
"""Core data models."""

from __future__ import annotations

import json
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any


# ── v4 Identity primitives ────────────────────────────────────────────────


class MemoryScope(str, Enum):
    """Visibility scope for a memory.

    USER     — visible across all sessions and projects of this user (default)
    PROJECT  — visible across all sessions in one project
    SESSION  — visible only within a single conversation session
    """
    USER = "user"
    PROJECT = "project"
    SESSION = "session"


def _now_utc() -> datetime:
    return datetime.now(timezone.utc)


@dataclass(frozen=True, slots=True)
class User:
    id: str
    external_id: str
    email: str | None = None
    display_name: str | None = None
    status: str = "active"
    created_at: datetime = field(default_factory=_now_utc)
    deleted_at: datetime | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_row(cls, row: dict[str, Any]) -> User:
        meta = row.get("metadata") or {}
        if isinstance(meta, str):
            meta = json.loads(meta)
        return cls(
            id=str(row["id"]),
            external_id=row["external_id"],
            email=row.get("email"),
            display_name=row.get("display_name"),
            status=row.get("status", "active"),
            created_at=row["created_at"],
            deleted_at=row.get("deleted_at"),
            metadata=meta,
        )


@dataclass(frozen=True, slots=True)
class Project:
    id: str
    user_id: str
    name: str
    description: str | None = None
    status: str = "active"
    created_at: datetime = field(default_factory=_now_utc)
    archived_at: datetime | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def is_inbox(self) -> bool:
        return bool(self.metadata.get("is_inbox"))

    @classmethod
    def from_row(cls, row: dict[str, Any]) -> Project:
        meta = row.get("metadata") or {}
        if isinstance(meta, str):
            meta = json.loads(meta)
        return cls(
            id=str(row["id"]),
            user_id=str(row["user_id"]),
            name=row["name"],
            description=row.get("description"),
            status=row.get("status", "active"),
            created_at=row["created_at"],
            archived_at=row.get("archived_at"),
            metadata=meta,
        )


@dataclass(frozen=True, slots=True)
class Session:
    id: str
    user_id: str
    project_id: str | None = None
    title: str | None = None
    status: str = "active"
    created_at: datetime = field(default_factory=_now_utc)
    last_active_at: datetime = field(default_factory=_now_utc)
    ended_at: datetime | None = None
    message_count: int = 0
    consolidation_state: str = "pending"
    metadata: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_row(cls, row: dict[str, Any]) -> Session:
        meta = row.get("metadata") or {}
        if isinstance(meta, str):
            meta = json.loads(meta)
        return cls(
            id=str(row["id"]),
            user_id=str(row["user_id"]),
            project_id=str(row["project_id"]) if row.get("project_id") else None,
            title=row.get("title"),
            status=row.get("status", "active"),
            created_at=row["created_at"],
            last_active_at=row["last_active_at"],
            ended_at=row.get("ended_at"),
            message_count=row.get("message_count", 0),
            consolidation_state=row.get("consolidation_state", "pending"),
            metadata=meta,
        )


# ── Existing Memory / RetrievalResult (Phase 0 keeps the v3 shape) ───────
# Phase 1 will extend Memory with v4 columns (user_id, scope, t_created, etc.)
# and rewrite namespace usage. For Phase 0 we ship the schema + identity
# primitives without touching the existing Memory dataclass — keeps every
# caller working until Phase 1 lands.


# ── Hierarchical memory layers (2026 best-practice: Letta / LangGraph / Mem0) ──
# Memory.layer classifies the *intent* of a memory, separate from category:
#
#   episodic   — what happened in conversation (default; conversational rounds)
#   semantic   — what is true (distilled facts; reflection / consolidation)
#   procedural — how to do things (skills, workflows, recipes)
#   working    — session-scoped, ephemeral (auto-purged on session end; the
#                cron path is out of scope for this PR — column + plumbing only)
#
# Layer is a CLOSED vocabulary — invalid values raise at the boundary so a
# typo can't silently land in the DB.
LAYER_EPISODIC: str = "episodic"
LAYER_SEMANTIC: str = "semantic"
LAYER_PROCEDURAL: str = "procedural"
LAYER_WORKING: str = "working"
LAYER_DEFAULT: str = LAYER_EPISODIC

VALID_LAYERS: frozenset[str] = frozenset({
    LAYER_EPISODIC, LAYER_SEMANTIC, LAYER_PROCEDURAL, LAYER_WORKING,
})


def _validate_layer(value: str) -> None:
    """Raise ValueError if ``value`` is not in :data:`VALID_LAYERS`.

    The error message lists every allowed layer so the caller doesn't have
    to look it up. Used at every system boundary that accepts a layer
    string (Memory.__post_init__, AgentMemory.add, AgentMemory.recall).
    """
    if value not in VALID_LAYERS:
        valid = ", ".join(sorted(VALID_LAYERS))
        raise ValueError(
            f"invalid layer {value!r}; must be one of: {valid}"
        )


@dataclass(frozen=True, slots=True)
class Memory:
    id: str = field(default_factory=lambda: uuid.uuid4().hex[:12])
    content: str = ""

    # Classification
    tags: list[str] = field(default_factory=list)
    category: str = "general"
    entity: str | None = None
    namespace: str = "default"   # legacy v3; derived from user/project/session in v4

    # Hierarchical memory layer (2026 best-practice).
    # See LAYER_* constants above. Default = 'episodic' so existing callers
    # that don't pass a layer keep their conversational-round behavior.
    # Validated in __post_init__ — ``Memory(layer='garbage')`` raises.
    layer: str = LAYER_DEFAULT

    # v4 tenancy (Optional → backward-compat with v3 callers)
    user_id: str | None = None
    project_id: str | None = None
    session_id: str | None = None
    scope: str = "user"          # one of: user | project | session

    # Temporal — event time
    created_at: str = field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )
    event_date: str | None = None
    valid_from: str = field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )
    valid_until: str | None = None
    superseded_by: str | None = None

    # v4 bi-temporal — transaction time (when system knew it)
    t_created: str | None = None
    t_expired: str | None = None

    # v4 provenance — links back to source episode + extraction context
    source_episode_id: str | None = None
    source_span: list[int] | None = None   # [start_char, end_char]
    extraction_model: str | None = None
    agent_id: str | None = None
    parent_agent_id: str | None = None
    visibility: str = "team"
    signature: str | None = None           # opt-in Ed25519 sig

    # Vector (optional)
    embedding: list[float] | None = None

    # Provenance
    confidence: float = 1.0

    # Status
    status: str = "active"

    # Access tracking
    access_count: int = 0
    last_accessed: str | None = None

    # Content dedup
    content_hash: str | None = None

    # Extensible
    metadata: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        # Closed-vocabulary check at the dataclass boundary — typos in
        # caller code (`layer="proceduiral"`) fail loudly here rather
        # than landing in the DB and getting filtered out at recall time.
        _validate_layer(self.layer)

    @classmethod
    def from_row(cls, row: dict[str, Any]) -> Memory:
        """Create a Memory from a document-store row dict.

        Backward-compatible with both v3 (TEXT id, namespace, no bi-temporal)
        and v4 (UUID id, user_id/scope, t_created/t_expired) rows. Missing
        columns get sensible defaults so v3 callers see no behavior change."""
        tags_raw = row["tags"]
        tags = json.loads(tags_raw) if isinstance(tags_raw, str) else tags_raw
        metadata_raw = row.get("metadata") or {}
        metadata = (
            json.loads(metadata_raw) if isinstance(metadata_raw, str) else metadata_raw
        )

        def _maybe_str(v: Any) -> str | None:
            return str(v) if v is not None else None

        def _maybe_iso(v: Any) -> str | None:
            if v is None:
                return None
            if isinstance(v, datetime):
                return v.isoformat()
            return str(v)

        # source_span comes back from psycopg2 as a Range object for INT4RANGE.
        span_raw = row.get("source_span")
        if span_raw is None:
            source_span = None
        elif hasattr(span_raw, "lower") and hasattr(span_raw, "upper"):
            # psycopg2 Range type
            source_span = [span_raw.lower, span_raw.upper]
        elif isinstance(span_raw, (list, tuple)):
            source_span = list(span_raw)
        else:
            source_span = None

        # created_at / valid_from may come back as datetime in v4 schema or
        # as ISO string in v3 schema. Always emit ISO string.
        created_at = _maybe_iso(row.get("created_at")) or datetime.now(timezone.utc).isoformat()
        valid_from = _maybe_iso(row.get("valid_from")) or created_at

        # Namespace resolution: v3 schema has a top-level `namespace` column;
        # v4 dropped it (tenancy moved to user_id + project_id + scope) so a
        # v4 read returns no namespace and the dataclass defaults to "default".
        # That breaks any caller that uses namespace as a metadata-style key
        # (e.g. the LME bench writes namespace="lme_<sample>" for sample
        # isolation across the shared DB). Fall back to metadata["_namespace"]
        # which AgentContext already stamps and which we now stamp from
        # core.add() too. Final fallback is the legacy "default".
        _ns_from_row = row.get("namespace")
        if _ns_from_row is None:
            _md_for_ns = metadata if isinstance(metadata, dict) else {}
            _ns_from_row = _md_for_ns.get("_namespace", "default")

        # Layer column was added in the hierarchical-layers PR. Pre-existing
        # rows (and v4 rows from before the migration) won't have it; default
        # them to ``episodic``. We also coerce unknown values to ``episodic``
        # rather than crashing on the read path — a corrupted DB row should
        # be inert, not poison every recall.
        _layer_from_row = row.get("layer")
        if not isinstance(_layer_from_row, str) or _layer_from_row not in VALID_LAYERS:
            _layer_from_row = LAYER_DEFAULT
        return cls(
            id=str(row["id"]),
            content=row["content"],
            tags=list(tags) if tags is not None else [],
            category=row.get("category", "general"),
            entity=row.get("entity"),
            namespace=_ns_from_row,
            layer=_layer_from_row,
            user_id=_maybe_str(row.get("user_id")),
            project_id=_maybe_str(row.get("project_id")),
            session_id=_maybe_str(row.get("session_id")),
            scope=row.get("scope", "user"),
            created_at=created_at,
            event_date=_maybe_iso(row.get("event_date")),
            valid_from=valid_from,
            valid_until=_maybe_iso(row.get("valid_until")),
            superseded_by=_maybe_str(row.get("superseded_by")),
            t_created=_maybe_iso(row.get("t_created")),
            t_expired=_maybe_iso(row.get("t_expired")),
            source_episode_id=_maybe_str(row.get("source_episode_id")),
            source_span=source_span,
            extraction_model=row.get("extraction_model"),
            agent_id=row.get("agent_id"),
            parent_agent_id=row.get("parent_agent_id"),
            visibility=row.get("visibility", "team"),
            signature=row.get("signature"),
            confidence=row.get("confidence", 1.0),
            status=row.get("status", "active"),
            access_count=row.get("access_count", 0),
            last_accessed=_maybe_iso(row.get("last_accessed")),
            content_hash=row.get("content_hash"),
            metadata=metadata,
        )

    def to_dict(self) -> dict[str, Any]:
        """Convert to a plain dict for JSON export.

        Includes both v3 and v4 fields. v4 fields are present but None when
        the memory was created via a v3 path."""
        return {
            "id": self.id,
            "content": self.content,
            "tags": self.tags,
            "category": self.category,
            "entity": self.entity,
            "namespace": self.namespace,
            "layer": self.layer,
            "user_id": self.user_id,
            "project_id": self.project_id,
            "session_id": self.session_id,
            "scope": self.scope,
            "created_at": self.created_at,
            "event_date": self.event_date,
            "valid_from": self.valid_from,
            "valid_until": self.valid_until,
            "superseded_by": self.superseded_by,
            "t_created": self.t_created,
            "t_expired": self.t_expired,
            "source_episode_id": self.source_episode_id,
            "source_span": self.source_span,
            "extraction_model": self.extraction_model,
            "agent_id": self.agent_id,
            "parent_agent_id": self.parent_agent_id,
            "visibility": self.visibility,
            "signature": self.signature,
            "confidence": self.confidence,
            "status": self.status,
            "access_count": self.access_count,
            "last_accessed": self.last_accessed,
            "content_hash": self.content_hash,
            "metadata": self.metadata,
        }

    @property
    def is_v4(self) -> bool:
        """True if this memory was written via the v4 tenancy path."""
        return self.user_id is not None


@dataclass(frozen=True, slots=True)
class RetrievalResult:
    memory: Memory
    score: float
    match_source: str  # "tag", "fts", "vector"

    @property
    def content(self) -> str:
        return self.memory.content


# ── ContextPack (Phase 6.1, roadmap §D.1) ────────────────────────────────


@dataclass(frozen=True, slots=True)
class ContextPackEntry:
    """One memory inside a ContextPack — citation-friendly view.

    The agent's Chain-of-Note instructions cite by ``id`` (e.g. [mem_42]),
    so the id MUST round-trip exactly. ``source_episode_id`` lets the
    auditor reconstruct the verbatim round the fact came from.
    """
    id: str
    content: str
    category: str
    entity: str | None
    valid_from: str | None
    valid_until: str | None
    confidence: float
    source_episode_id: str | None
    score: float

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "content": self.content,
            "category": self.category,
            "entity": self.entity,
            "valid_from": self.valid_from,
            "valid_until": self.valid_until,
            "confidence": self.confidence,
            "source_episode_id": self.source_episode_id,
            "score": self.score,
        }


@dataclass(frozen=True, slots=True)
class ContextPack:
    """Structured retrieval envelope for Chain-of-Note consumption.

    Returned by ``AgentMemory.recall_as_context``. The agent prepends
    ``chain_of_note_prompt`` to its system prompt, then renders
    ``memories`` (already sorted by score) into the {memories_json}
    placeholder.

    Why structured: the agent needs to cite ids, abstain when irrelevant,
    and prefer the right validity window when memories conflict — all
    impossible if recall_as_context returns plain text.
    """
    query: str
    memories: list[ContextPackEntry]
    as_of: str | None
    token_count: int
    chain_of_note_prompt: str

    @property
    def memory_count(self) -> int:
        return len(self.memories)

    def to_dict(self) -> dict[str, Any]:
        return {
            "query": self.query,
            "memories": [m.to_dict() for m in self.memories],
            "as_of": self.as_of,
            "token_count": self.token_count,
            "chain_of_note_prompt": self.chain_of_note_prompt,
        }

    def memories_json(self) -> str:
        """Render the memories list as a JSON string suitable for the
        {memories_json} placeholder in the Chain-of-Note prompt."""
        return json.dumps([m.to_dict() for m in self.memories],
                          ensure_ascii=False, indent=2)

    def render_prompt(self) -> str:
        """Return the Chain-of-Note prompt with memories interpolated."""
        return self.chain_of_note_prompt.format(
            memories_json=self.memories_json(),
        )


