"""Add memory tables and user encryption_key column.

Memory tables:
  memory_core        — per-user key/value preferences (encrypted)
  memory_associative — semantic memory with pgvector embedding (encrypted)
  memory_episodic    — session summaries (encrypted)
  memory_proactive   — behavioral patterns (encrypted)

Also adds encryption_key column to users table.

Revision ID: 004
Revises: 003
Create Date: 2026-03-08
"""

from __future__ import annotations

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql

revision: str = "004"
down_revision: Union[str, None] = "003"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ── Enable pgvector extension (idempotent) ────────────────────────────────
    op.execute("CREATE EXTENSION IF NOT EXISTS vector;")

    # ── Add encryption_key to users ───────────────────────────────────────────
    op.add_column(
        "users",
        sa.Column("encryption_key", sa.String(64), nullable=True),
    )

    # ── memory_core ───────────────────────────────────────────────────────────
    op.create_table(
        "memory_core",
        sa.Column("id", postgresql.UUID(as_uuid=False), primary_key=True),
        sa.Column(
            "user_id",
            postgresql.UUID(as_uuid=False),
            sa.ForeignKey("users.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column("key", sa.String(255), nullable=False),
        sa.Column("value_encrypted", sa.Text, nullable=False),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
    )
    op.create_index("ix_memory_core_user_id", "memory_core", ["user_id"])

    # ── memory_associative ────────────────────────────────────────────────────
    # The embedding column uses pgvector's vector(1536) type.
    op.create_table(
        "memory_associative",
        sa.Column("id", postgresql.UUID(as_uuid=False), primary_key=True),
        sa.Column(
            "user_id",
            postgresql.UUID(as_uuid=False),
            sa.ForeignKey("users.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column("content_encrypted", sa.Text, nullable=False),
        sa.Column("entity_type", sa.String(100), nullable=True),
        sa.Column("entity_id", sa.String(255), nullable=True),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
    )
    # Add the pgvector column separately (not supported by generic sa types)
    op.execute(
        "ALTER TABLE memory_associative ADD COLUMN embedding vector(1536);"
    )
    op.create_index("ix_memory_associative_user_id", "memory_associative", ["user_id"])
    # IVFFlat index for approximate nearest-neighbour search
    op.execute(
        "CREATE INDEX ix_memory_associative_embedding "
        "ON memory_associative USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);"
    )

    # ── memory_episodic ───────────────────────────────────────────────────────
    op.create_table(
        "memory_episodic",
        sa.Column("id", postgresql.UUID(as_uuid=False), primary_key=True),
        sa.Column(
            "user_id",
            postgresql.UUID(as_uuid=False),
            sa.ForeignKey("users.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column("summary_encrypted", sa.Text, nullable=False),
        sa.Column("session_id", sa.String(255), nullable=False),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
    )
    op.create_index("ix_memory_episodic_user_id", "memory_episodic", ["user_id"])
    op.create_index("ix_memory_episodic_session_id", "memory_episodic", ["session_id"])

    # ── memory_proactive ──────────────────────────────────────────────────────
    op.create_table(
        "memory_proactive",
        sa.Column("id", postgresql.UUID(as_uuid=False), primary_key=True),
        sa.Column(
            "user_id",
            postgresql.UUID(as_uuid=False),
            sa.ForeignKey("users.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column("pattern_encrypted", sa.Text, nullable=False),
        sa.Column("confidence", sa.Float, nullable=False, server_default="0.5"),
        sa.Column("source", sa.String(50), nullable=False, server_default="inferred"),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
    )
    op.create_index("ix_memory_proactive_user_id", "memory_proactive", ["user_id"])


def downgrade() -> None:
    op.drop_table("memory_proactive")
    op.drop_table("memory_episodic")
    op.drop_index("ix_memory_associative_embedding", "memory_associative")
    op.drop_table("memory_associative")
    op.drop_table("memory_core")
    op.drop_column("users", "encryption_key")