"""Thread-scoped conversation memory backed by SQLAlchemy async."""

from __future__ import annotations

import uuid
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any, Literal

try:
    from sqlalchemy import Column, DateTime, String, Text, select
    from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
    from sqlalchemy.orm import DeclarativeBase, sessionmaker
except ImportError as e:
    raise ImportError("SQLAlchemy not installed. Install with: pip install koan[memory]") from e

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncEngine

from koan.memory.types import ThreadMessage


class _Base(DeclarativeBase):
    pass


class _MessageRow(_Base):
    __tablename__ = "thread_messages"

    id = Column(String, primary_key=True)
    thread_id = Column(String, nullable=False, index=True)
    user_id = Column(String, nullable=True, index=True)
    role = Column(String, nullable=False)
    content = Column(Text, nullable=False)
    name = Column(String, nullable=True)
    metadata_json = Column(Text, nullable=False, default="{}")
    created_at = Column(DateTime, nullable=False, default=datetime.now(UTC))
    sequence = Column(String, nullable=False)


class ThreadStore:
    """Thread-scoped conversation memory.

    Stores messages per thread_id using SQLAlchemy async with SQLite
    (or any async-compatible database).

    Usage::

        store = ThreadStore("sqlite+aiosqlite:///memory.db")
        await store.init()

        await store.append("thread-1", role="user", content="Hello")
        messages = await store.get_messages("thread-1")
    """

    def __init__(self, db_url: str = "sqlite+aiosqlite://") -> None:
        self._engine: AsyncEngine = create_async_engine(db_url)
        self._session_factory = sessionmaker(self._engine, class_=AsyncSession, expire_on_commit=False)  # type: ignore[call-overload]

    async def init(self) -> None:
        """Create tables if they don't exist."""
        async with self._engine.begin() as conn:
            await conn.run_sync(_Base.metadata.create_all)

    async def append(
        self,
        thread_id: str,
        *,
        role: Literal["system", "user", "assistant", "tool"],
        content: str,
        name: str | None = None,
        user_id: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> ThreadMessage:
        """Append a message to a thread."""
        import orjson

        msg_id = str(uuid.uuid4())
        now = datetime.now(UTC)
        seq = now.isoformat() + "-" + msg_id[:8]

        row = _MessageRow(
            id=msg_id,
            thread_id=thread_id,
            user_id=user_id,
            role=role,
            content=content,
            name=name,
            metadata_json=orjson.dumps(metadata or {}).decode(),
            created_at=now,
            sequence=seq,
        )

        async with self._session_factory() as session:  # type: ignore[attr-defined]
            session.add(row)
            await session.commit()

        return ThreadMessage(
            id=msg_id,
            thread_id=thread_id,
            role=role,
            content=content,
            name=name,
            metadata=metadata or {},
            created_at=now,
        )

    async def get_messages(self, thread_id: str, *, limit: int | None = None) -> list[ThreadMessage]:
        """Get all messages in a thread, ordered by sequence."""
        import orjson

        stmt = (
            select(_MessageRow)
            .where(_MessageRow.thread_id == thread_id)
            .order_by(_MessageRow.sequence)
        )
        if limit is not None:
            stmt = stmt.limit(limit)

        async with self._session_factory() as session:  # type: ignore[attr-defined]
            result = await session.execute(stmt)
            rows = result.scalars().all()

        return [
            ThreadMessage(
                id=row.id,
                thread_id=row.thread_id,
                role=row.role,
                content=row.content,
                name=row.name,
                metadata=orjson.loads(row.metadata_json) if row.metadata_json else {},
                created_at=row.created_at,
            )
            for row in rows
        ]

    async def list_threads(self, *, user_id: str | None = None) -> list[str]:
        """List distinct thread IDs, optionally filtered by user_id."""
        from sqlalchemy import distinct

        stmt = select(distinct(_MessageRow.thread_id))
        if user_id is not None:
            stmt = stmt.where(_MessageRow.user_id == user_id)
        stmt = stmt.order_by(_MessageRow.thread_id)

        async with self._session_factory() as session:  # type: ignore[attr-defined]
            result = await session.execute(stmt)
            return [row[0] for row in result.all()]

    async def delete_thread(self, thread_id: str) -> int:
        """Delete all messages in a thread. Returns count of deleted messages."""
        from sqlalchemy import delete

        stmt = delete(_MessageRow).where(_MessageRow.thread_id == thread_id)
        async with self._session_factory() as session:  # type: ignore[attr-defined]
            result = await session.execute(stmt)
            await session.commit()
            return result.rowcount  # type: ignore[return-value]

    async def close(self) -> None:
        """Dispose of the engine."""
        await self._engine.dispose()
