"""Durable local graph-memory schema replacing the storage side of Zep."""

from __future__ import annotations

from datetime import datetime, timezone
from uuid import uuid4

from sqlalchemy import DateTime, Float, ForeignKey, Index, JSON, String, Text, UniqueConstraint, func, text
from sqlalchemy.orm import Mapped, mapped_column, relationship

from ..db import Base


def _id(prefix: str) -> str:
    return f"{prefix}_{uuid4().hex}"


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


class MemoryGraph(Base):
    __tablename__ = "memory_graphs"

    id: Mapped[str] = mapped_column(String(128), primary_key=True, default=lambda: _id("graph"))
    organization_id: Mapped[str] = mapped_column(String(64), nullable=False, index=True)
    project_id: Mapped[str] = mapped_column(String(128), nullable=False, index=True)
    ontology: Mapped[dict] = mapped_column(JSON, nullable=False, default=dict, server_default=text("'{}'"))
    status: Mapped[str] = mapped_column(String(32), nullable=False, default="ready", server_default=text("'ready'"))
    version: Mapped[int] = mapped_column(nullable=False, default=1, server_default=text("1"))
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=_utc_now, server_default=func.now()
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=_utc_now, server_default=func.now(), onupdate=_utc_now
    )

    episodes: Mapped[list["MemoryEpisode"]] = relationship(
        back_populates="graph", cascade="all, delete-orphan"
    )
    nodes: Mapped[list["MemoryNode"]] = relationship(
        back_populates="graph", cascade="all, delete-orphan"
    )
    edges: Mapped[list["MemoryEdge"]] = relationship(
        back_populates="graph", cascade="all, delete-orphan"
    )


class MemoryEpisode(Base):
    __tablename__ = "memory_episodes"
    __table_args__ = (
        UniqueConstraint("graph_id", "source_type", "source_ref", name="uq_memory_episode_source"),
        Index("ix_memory_episodes_graph_status", "graph_id", "status"),
    )

    id: Mapped[str] = mapped_column(String(128), primary_key=True, default=lambda: _id("episode"))
    graph_id: Mapped[str] = mapped_column(
        ForeignKey("memory_graphs.id", ondelete="CASCADE"), nullable=False, index=True
    )
    source_type: Mapped[str] = mapped_column(String(32), nullable=False)
    source_ref: Mapped[str] = mapped_column(String(256), nullable=False)
    normalized_text: Mapped[str] = mapped_column(Text, nullable=False)
    summary: Mapped[str] = mapped_column(Text, nullable=False, default="", server_default=text("''"))
    status: Mapped[str] = mapped_column(String(32), nullable=False, default="processed", server_default=text("'processed'"))
    extractor_version: Mapped[str] = mapped_column(String(64), nullable=False, default="v1", server_default=text("'v1'"))
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=_utc_now, server_default=func.now()
    )

    graph: Mapped[MemoryGraph] = relationship(back_populates="episodes")


class MemoryNode(Base):
    __tablename__ = "memory_nodes"
    __table_args__ = (
        UniqueConstraint("graph_id", "normalized_name", name="uq_memory_node_graph_name"),
        Index("ix_memory_nodes_graph_name", "graph_id", "normalized_name"),
    )

    id: Mapped[str] = mapped_column(String(128), primary_key=True, default=lambda: _id("node"))
    graph_id: Mapped[str] = mapped_column(
        ForeignKey("memory_graphs.id", ondelete="CASCADE"), nullable=False, index=True
    )
    canonical_name: Mapped[str] = mapped_column(String(512), nullable=False)
    normalized_name: Mapped[str] = mapped_column(String(512), nullable=False)
    labels: Mapped[list] = mapped_column(JSON, nullable=False, default=list, server_default=text("'[]'"))
    aliases: Mapped[list] = mapped_column(JSON, nullable=False, default=list, server_default=text("'[]'"))
    attributes: Mapped[dict] = mapped_column(JSON, nullable=False, default=dict, server_default=text("'{}'"))
    summary: Mapped[str] = mapped_column(Text, nullable=False, default="", server_default=text("''"))
    confidence: Mapped[float] = mapped_column(Float, nullable=False, default=0.0, server_default=text("0"))
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=_utc_now, server_default=func.now()
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=_utc_now, server_default=func.now(), onupdate=_utc_now
    )

    graph: Mapped[MemoryGraph] = relationship(back_populates="nodes")
    outgoing_edges: Mapped[list["MemoryEdge"]] = relationship(
        foreign_keys="MemoryEdge.source_node_id", back_populates="source_node"
    )
    incoming_edges: Mapped[list["MemoryEdge"]] = relationship(
        foreign_keys="MemoryEdge.target_node_id", back_populates="target_node"
    )


class MemoryEdge(Base):
    __tablename__ = "memory_edges"
    __table_args__ = (
        Index("ix_memory_edges_graph_relation", "graph_id", "relation"),
        Index("ix_memory_edges_graph_temporal", "graph_id", "valid_at", "invalid_at"),
    )

    id: Mapped[str] = mapped_column(String(128), primary_key=True, default=lambda: _id("edge"))
    graph_id: Mapped[str] = mapped_column(
        ForeignKey("memory_graphs.id", ondelete="CASCADE"), nullable=False, index=True
    )
    source_node_id: Mapped[str] = mapped_column(
        ForeignKey("memory_nodes.id", ondelete="CASCADE"), nullable=False, index=True
    )
    target_node_id: Mapped[str] = mapped_column(
        ForeignKey("memory_nodes.id", ondelete="CASCADE"), nullable=False, index=True
    )
    relation: Mapped[str] = mapped_column(String(128), nullable=False)
    fact: Mapped[str] = mapped_column(Text, nullable=False)
    attributes: Mapped[dict] = mapped_column(JSON, nullable=False, default=dict, server_default=text("'{}'"))
    confidence: Mapped[float] = mapped_column(Float, nullable=False, default=0.0, server_default=text("0"))
    valid_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    invalid_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    expired_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=_utc_now, server_default=func.now()
    )

    graph: Mapped[MemoryGraph] = relationship(back_populates="edges")
    source_node: Mapped[MemoryNode] = relationship(
        foreign_keys=[source_node_id], back_populates="outgoing_edges"
    )
    target_node: Mapped[MemoryNode] = relationship(
        foreign_keys=[target_node_id], back_populates="incoming_edges"
    )