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

import uuid
from datetime import datetime, timezone

from sqlalchemy import Column, DateTime, Float, String, Text
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


class ExperienceModel(Base):
    """Task experience ORM model for RiskGuardLearner / ExperienceStore.

    Stores task execution outcomes (success/failure/partial) with optional
    pgvector embeddings for semantic similarity search.
    """

    __tablename__ = "task_experiences"

    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    task_type = Column(String, index=True)
    goal = Column(Text, default="")
    steps_summary = Column(Text, default="")
    outcome = Column(String, default="success")  # "success", "failure", "partial"
    duration_seconds = Column(Float, default=0.0)
    success_rate = Column(Float, default=1.0)
    failure_reasons = Column(JSONB, default=list)  # list[str]
    optimization_tips = Column(JSONB, default=list)  # list[str]
    embedding = Column(Text, nullable=True)  # JSON-encoded float list
    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), index=True)


def create_experience_session_factory(database_url: str):
    """Create an async session factory for task experiences.

    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()