"""SQLAlchemy ORM models for episodic memory persistence (PostgreSQL + pgvector)."""

import uuid
from datetime import datetime, timezone

from sqlalchemy import Column, DateTime, Float, String, Text, create_engine
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import declarative_base, sessionmaker

Base = declarative_base()


class EpisodeModel(Base):
    """Episodic memory ORM model

    Stores task execution experiences with optional pgvector embeddings
    for semantic similarity search.
    """

    __tablename__ = "episodic_memories"

    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    agent_name = Column(String, index=True)
    task_type = Column(String, index=True)
    input_summary = Column(Text, default="")
    output_summary = Column(Text, default="")
    outcome = Column(String, default="success")  # "success", "failure", "partial"
    quality_score = Column(Float, default=0.5)
    reflection = Column(Text, default="")
    embedding = Column(Text, nullable=True)  # JSON-encoded float list; pgvector if extension available
    metadata_ = Column("metadata", JSONB, nullable=True)  # Additional metadata
    created_at = Column(
        DateTime, default=lambda: datetime.now(timezone.utc), index=True
    )


def create_episodic_session_factory(database_url: str):
    """Create an async session factory for episodic memory.

    Args:
        database_url: PostgreSQL connection string,
            e.g. "postgresql+asyncpg://user:pass@localhost/dbname"

    Returns:
        async_sessionmaker bound to the engine.
    """
    from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine

    engine = create_async_engine(database_url, echo=False)
    async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
    return async_session


async def ensure_episodic_table(database_url: str) -> None:
    """Create the episodic_memories table if it does not exist.

    Safe to call on startup — uses CREATE TABLE IF NOT EXISTS.
    """
    from sqlalchemy.ext.asyncio import create_async_engine

    engine = create_async_engine(database_url, echo=False)
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    await engine.dispose()
