from datetime import datetime
from typing import Any

from sqlalchemy import CursorResult, Float, Select, cast, delete, func, literal, or_, select
from sqlalchemy.dialects.postgresql import REGCONFIG
from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)

from session_hub_core.config import DatabaseConfig
from session_hub_core.embeddings import EmbeddingService
from session_hub_core.exceptions import QueryError, SessionNotFoundError
from session_hub_core.models import (
    HostInfo,
    HostStats,
    MessageModel,
    MessageRow,
    MessageSearchResult,
    SessionDetail,
    SessionModel,
    SessionRow,
    SessionSearchResult,
    Stats,
    SyncLogModel,
    ToolStats,
)


class QueryAPI:
    """Query the central session database."""

    def __init__(self, config: DatabaseConfig, embeddings: EmbeddingService | None = None) -> None:
        self._config = config
        self._embeddings = embeddings
        self._engine: AsyncEngine = create_async_engine(
            config.dsn,
            pool_size=config.pool_size,
            max_overflow=config.max_overflow,
            pool_pre_ping=True,
            pool_recycle=300,
            echo=config.echo,
        )
        self._session_factory: async_sessionmaker[AsyncSession] = async_sessionmaker(
            self._engine, expire_on_commit=False
        )

    async def list_sessions(
        self,
        tool: str | None = None,
        host: str | None = None,
        working_directory: str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[SessionRow]:
        """List sessions with optional filters.

        Args:
            tool: Filter by tool name.
            host: Filter by host.
            working_directory: Substring match on working directory.
            limit: Maximum results (default 50).
            offset: Pagination offset.

        Returns:
            List of SessionRow, most recent activity first.
        """
        stmt = (
            select(SessionModel)
            .order_by(SessionModel.last_activity_at.desc())
            .limit(limit)
            .offset(offset)
        )
        if tool is not None:
            stmt = stmt.where(SessionModel.tool == tool)
        if host is not None:
            stmt = stmt.where(SessionModel.host == host)
        if working_directory is not None:
            stmt = stmt.where(SessionModel.working_directory.ilike(f"%{working_directory}%"))
        async with self._session_factory() as conn:
            rows = (await conn.execute(stmt)).scalars().all()
        return [self._session_row(row) for row in rows]

    async def get_session(
        self,
        host: str,
        tool: str,
        session_id: str,
    ) -> SessionDetail:
        """Get a single session with all messages.

        Args:
            host: Host that synced the session.
            tool: Tool name.
            session_id: Session identifier.

        Raises:
            SessionNotFoundError: If no session matches.
        """
        async with self._session_factory() as conn:
            session = (
                await conn.execute(
                    select(SessionModel).where(
                        SessionModel.host == host,
                        SessionModel.tool == tool,
                        SessionModel.session_id == session_id,
                    )
                )
            ).scalar_one_or_none()
            if session is None:
                raise SessionNotFoundError(f"Session {host}/{tool}/{session_id} not found")
            messages = (
                (
                    await conn.execute(
                        select(MessageModel)
                        .where(MessageModel.session_pk == session.id)
                        .order_by(MessageModel.id)
                    )
                )
                .scalars()
                .all()
            )
        return SessionDetail(
            **self._session_row(session).model_dump(),
            messages=[self._message_row(m) for m in messages],
            metadata=session.metadata_,
        )

    async def search_messages(
        self,
        query: str,
        tool: str | None = None,
        host: str | None = None,
        role: str | None = None,
        limit: int = 50,
        mode: str = "fts",
    ) -> list[MessageSearchResult]:
        """Search across all message content.

        Args:
            query: Search query.
            tool: Filter by tool.
            host: Filter by host.
            role: Filter by message role (user, assistant, etc.).
            limit: Maximum results.
            mode: Search mode — "fts" (full-text, default), "fuzzy"
                  (trigram, typo-tolerant), "semantic" (pgvector),
                  "hybrid" (FTS + semantic, combined ranking).

        Returns:
            List of MessageSearchResult ranked by relevance.
        """
        if mode == "fts":
            return await self._search_messages_fts(query, tool, host, role, limit)
        if mode == "fuzzy":
            return await self._search_messages_fuzzy(query, tool, host, role, limit)
        if mode == "semantic":
            return await self._search_messages_semantic(query, tool, host, role, limit)
        if mode == "hybrid":
            return await self._search_messages_hybrid(query, tool, host, role, limit)
        raise QueryError(f"Unknown search mode: {mode}")

    async def search_sessions(
        self,
        query: str,
        tool: str | None = None,
        host: str | None = None,
        limit: int = 50,
        mode: str = "fts",
    ) -> list[SessionRow]:
        """Search across session titles and content.

        Args:
            query: Search query.
            tool: Filter by tool.
            host: Filter by host.
            limit: Maximum results.
            mode: Search mode — "fts", "fuzzy", "semantic", "hybrid".

        Returns:
            Sessions whose title or content matches, ranked by relevance.
        """
        if mode == "fts":
            return await self._search_sessions_fts(query, tool, host, limit)
        if mode == "fuzzy":
            return await self._search_sessions_fuzzy(query, tool, host, limit)
        if mode == "semantic":
            return await self._search_sessions_semantic(query, tool, host, limit)
        if mode == "hybrid":
            return await self._search_sessions_hybrid(query, tool, host, limit)
        raise QueryError(f"Unknown search mode: {mode}")

    async def search_similar_sessions(
        self,
        query: str,
        tool: str | None = None,
        host: str | None = None,
        limit: int = 20,
    ) -> list[SessionSearchResult]:
        """Semantic search via pgvector.

        Embeds the query via OpenRouter, finds nearest sessions by cosine
        distance using the HNSW index. Returns sessions with similarity
        scores.

        Requires OPENROUTER_API_KEY to be configured.
        """
        embeddings = self._require_embeddings()
        return await embeddings.search_similar_sessions(query, tool=tool, host=host, limit=limit)

    async def search_similar_messages(
        self,
        query: str,
        tool: str | None = None,
        host: str | None = None,
        role: str | None = None,
        limit: int = 50,
    ) -> list[MessageSearchResult]:
        """Semantic search via pgvector on message-level embeddings.

        Embeds the query, finds nearest messages by cosine distance.
        """
        embeddings = self._require_embeddings()
        return await embeddings.search_similar_messages(
            query, tool=tool, host=host, role=role, limit=limit
        )

    async def get_stats(self) -> Stats:
        """Get aggregate statistics.

        Returns counts by host, tool, and time period.
        """
        async with self._session_factory() as conn:
            total_sessions = (
                await conn.execute(select(func.count()).select_from(SessionModel))
            ).scalar_one()
            total_messages = (
                await conn.execute(select(func.count()).select_from(MessageModel))
            ).scalar_one()
            session_hosts = (
                await conn.execute(
                    select(SessionModel.host, func.count()).group_by(SessionModel.host)
                )
            ).all()
            message_hosts = (
                await conn.execute(
                    select(MessageModel.host, func.count()).group_by(MessageModel.host)
                )
            ).all()
            session_tools = (
                await conn.execute(
                    select(SessionModel.tool, func.count()).group_by(SessionModel.tool)
                )
            ).all()
            message_tools = (
                await conn.execute(
                    select(MessageModel.tool, func.count()).group_by(MessageModel.tool)
                )
            ).all()
            last_syncs: dict[str, datetime | None] = {}
            for sync_host, started_at in (
                await conn.execute(
                    select(SyncLogModel.host, func.max(SyncLogModel.started_at))
                    .where(SyncLogModel.status == "success")
                    .group_by(SyncLogModel.host)
                )
            ).all():
                last_syncs[sync_host] = started_at
            oldest, newest = (
                await conn.execute(
                    select(func.min(SessionModel.created_at), func.max(SessionModel.created_at))
                )
            ).one()
        message_counts_by_host: dict[str, int] = {}
        for message_host, count in message_hosts:
            message_counts_by_host[message_host] = count
        message_counts_by_tool: dict[str, int] = {}
        for message_tool, count in message_tools:
            message_counts_by_tool[message_tool] = count
        by_host: dict[str, HostStats] = {
            host: HostStats(
                sessions=count,
                messages=message_counts_by_host.get(host, 0),
                last_sync=last_syncs.get(host),
            )
            for host, count in session_hosts
        }
        by_tool: dict[str, ToolStats] = {
            tool: ToolStats(sessions=count, messages=message_counts_by_tool.get(tool, 0))
            for tool, count in session_tools
        }
        return Stats(
            total_sessions=total_sessions,
            total_messages=total_messages,
            by_host=by_host,
            by_tool=by_tool,
            oldest_session=oldest,
            newest_session=newest,
        )

    async def get_hosts(self) -> list[HostInfo]:
        """List all hosts that have synced, with last sync time."""
        async with self._session_factory() as conn:
            host_rows = (
                await conn.execute(
                    select(
                        SessionModel.host,
                        func.count(func.distinct(SessionModel.id)),
                        func.count(MessageModel.id),
                    )
                    .outerjoin(MessageModel, MessageModel.session_pk == SessionModel.id)
                    .group_by(SessionModel.host)
                )
            ).all()
            last_syncs: dict[str, datetime | None] = {}
            for sync_host, started_at in (
                await conn.execute(
                    select(SyncLogModel.host, func.max(SyncLogModel.started_at))
                    .where(SyncLogModel.status == "success")
                    .group_by(SyncLogModel.host)
                )
            ).all():
                last_syncs[sync_host] = started_at
        merged: dict[str, tuple[int, int]] = {
            host: (sessions, messages) for host, sessions, messages in host_rows
        }
        hosts = set(merged) | set(last_syncs)
        return [
            HostInfo(
                host=host,
                last_sync=last_syncs.get(host),
                session_count=merged.get(host, (0, 0))[0],
                message_count=merged.get(host, (0, 0))[1],
            )
            for host in sorted(hosts)
        ]

    async def delete_host_data(self, host: str) -> int:
        """Delete all sessions and messages for a host.

        Returns the number of sessions deleted.
        """
        async with self._session_factory() as conn:
            result = await conn.execute(delete(SessionModel).where(SessionModel.host == host))
            await conn.commit()
        assert isinstance(result, CursorResult)
        return result.rowcount or 0

    async def close(self) -> None:
        await self._engine.dispose()

    async def _search_messages_fts(
        self,
        query: str,
        tool: str | None,
        host: str | None,
        role: str | None,
        limit: int,
    ) -> list[MessageSearchResult]:
        tsq = func.plainto_tsquery(cast("english", REGCONFIG), func.unaccent(query))
        rank = func.ts_rank_cd(MessageModel.content_tsv, tsq)
        stmt = (
            select(
                MessageModel,
                SessionModel.title,
                SessionModel.host,
                SessionModel.tool,
                SessionModel.session_id,
                SessionModel.working_directory,
                rank.label("rank"),
            )
            .join(SessionModel, SessionModel.id == MessageModel.session_pk)
            .where(MessageModel.content_tsv.op("@@")(tsq))
            .order_by(rank.desc())
            .limit(limit)
        )
        stmt = self._apply_message_filters(stmt, tool, host, role)
        return await self._execute_message_search(stmt)

    async def _search_messages_fuzzy(
        self,
        query: str,
        tool: str | None,
        host: str | None,
        role: str | None,
        limit: int,
    ) -> list[MessageSearchResult]:
        q = func.unaccent(query)
        rank = func.greatest(
            func.similarity(MessageModel.content, q),
            func.word_similarity(q, MessageModel.content),
        )
        stmt = (
            select(
                MessageModel,
                SessionModel.title,
                SessionModel.host,
                SessionModel.tool,
                SessionModel.session_id,
                SessionModel.working_directory,
                rank.label("rank"),
            )
            .join(SessionModel, SessionModel.id == MessageModel.session_pk)
            .where(
                or_(
                    func.similarity(MessageModel.content, q) >= 0.3,
                    func.word_similarity(q, MessageModel.content) >= 0.3,
                )
            )
            .order_by(rank.desc())
            .limit(limit)
        )
        stmt = self._apply_message_filters(stmt, tool, host, role)
        return await self._execute_message_search(stmt)

    async def _search_messages_semantic(
        self,
        query: str,
        tool: str | None,
        host: str | None,
        role: str | None,
        limit: int,
    ) -> list[MessageSearchResult]:
        embeddings = self._require_embeddings()
        vector = (await embeddings.embed_batch([query]))[0]
        distance = MessageModel.embedding.op("<=>")(vector)
        stmt = (
            select(
                MessageModel,
                SessionModel.title,
                SessionModel.host,
                SessionModel.tool,
                SessionModel.session_id,
                SessionModel.working_directory,
                (literal(1.0, Float) - distance).label("rank"),
            )
            .join(SessionModel, SessionModel.id == MessageModel.session_pk)
            .where(MessageModel.embedding.is_not(None))
            .order_by(distance)
            .limit(limit)
        )
        stmt = self._apply_message_filters(stmt, tool, host, role)
        return await self._execute_message_search(stmt)

    async def _search_messages_hybrid(
        self,
        query: str,
        tool: str | None,
        host: str | None,
        role: str | None,
        limit: int,
    ) -> list[MessageSearchResult]:
        embeddings = self._require_embeddings()
        vector = (await embeddings.embed_batch([query]))[0]
        tsq = func.plainto_tsquery(cast("english", REGCONFIG), func.unaccent(query))
        ts_rank = func.ts_rank_cd(MessageModel.content_tsv, tsq)
        distance = MessageModel.embedding.op("<=>")(vector)
        score = 0.7 * ts_rank + 0.3 * (literal(1.0, Float) - distance)
        stmt = (
            select(
                MessageModel,
                SessionModel.title,
                SessionModel.host,
                SessionModel.tool,
                SessionModel.session_id,
                SessionModel.working_directory,
                score.label("rank"),
            )
            .join(SessionModel, SessionModel.id == MessageModel.session_pk)
            .where(
                MessageModel.content_tsv.op("@@")(tsq),
                MessageModel.embedding.is_not(None),
            )
            .order_by(score.desc())
            .limit(limit)
        )
        stmt = self._apply_message_filters(stmt, tool, host, role)
        return await self._execute_message_search(stmt)

    async def _search_sessions_fts(
        self,
        query: str,
        tool: str | None,
        host: str | None,
        limit: int,
    ) -> list[SessionRow]:
        tsq = func.plainto_tsquery(cast("english", REGCONFIG), func.unaccent(query))
        rank = func.ts_rank_cd(SessionModel.content_tsv, tsq)
        stmt = (
            select(SessionModel)
            .where(SessionModel.content_tsv.op("@@")(tsq))
            .order_by(rank.desc())
            .limit(limit)
        )
        stmt = self._apply_session_filters(stmt, tool, host)
        return await self._execute_session_search(stmt)

    async def _search_sessions_fuzzy(
        self,
        query: str,
        tool: str | None,
        host: str | None,
        limit: int,
    ) -> list[SessionRow]:
        q = func.unaccent(query)
        rank = func.greatest(
            func.similarity(SessionModel.title, q),
            func.word_similarity(q, SessionModel.title),
        )
        stmt = (
            select(SessionModel)
            .where(
                or_(
                    func.similarity(SessionModel.title, q) >= 0.3,
                    func.word_similarity(q, SessionModel.title) >= 0.3,
                )
            )
            .order_by(rank.desc())
            .limit(limit)
        )
        stmt = self._apply_session_filters(stmt, tool, host)
        return await self._execute_session_search(stmt)

    async def _search_sessions_semantic(
        self,
        query: str,
        tool: str | None,
        host: str | None,
        limit: int,
    ) -> list[SessionRow]:
        embeddings = self._require_embeddings()
        vector = (await embeddings.embed_batch([query]))[0]
        distance = SessionModel.embedding.op("<=>")(vector)
        stmt = (
            select(SessionModel)
            .where(SessionModel.embedding.is_not(None))
            .order_by(distance)
            .limit(limit)
        )
        stmt = self._apply_session_filters(stmt, tool, host)
        return await self._execute_session_search(stmt)

    async def _search_sessions_hybrid(
        self,
        query: str,
        tool: str | None,
        host: str | None,
        limit: int,
    ) -> list[SessionRow]:
        embeddings = self._require_embeddings()
        vector = (await embeddings.embed_batch([query]))[0]
        tsq = func.plainto_tsquery(cast("english", REGCONFIG), func.unaccent(query))
        ts_rank = func.ts_rank_cd(SessionModel.content_tsv, tsq)
        distance = SessionModel.embedding.op("<=>")(vector)
        score = 0.7 * ts_rank + 0.3 * (literal(1.0, Float) - distance)
        stmt = (
            select(SessionModel)
            .where(
                SessionModel.content_tsv.op("@@")(tsq),
                SessionModel.embedding.is_not(None),
            )
            .order_by(score.desc())
            .limit(limit)
        )
        stmt = self._apply_session_filters(stmt, tool, host)
        return await self._execute_session_search(stmt)

    def _require_embeddings(self) -> EmbeddingService:
        if self._embeddings is None:
            raise QueryError("Semantic search requires an EmbeddingService")
        return self._embeddings

    def _apply_message_filters(
        self,
        stmt: Select[Any],
        tool: str | None,
        host: str | None,
        role: str | None,
    ) -> Select[Any]:
        if tool is not None:
            stmt = stmt.where(MessageModel.tool == tool)
        if host is not None:
            stmt = stmt.where(MessageModel.host == host)
        if role is not None:
            stmt = stmt.where(MessageModel.role == role)
        return stmt

    def _apply_session_filters(
        self,
        stmt: Select[Any],
        tool: str | None,
        host: str | None,
    ) -> Select[Any]:
        if tool is not None:
            stmt = stmt.where(SessionModel.tool == tool)
        if host is not None:
            stmt = stmt.where(SessionModel.host == host)
        return stmt

    async def _execute_message_search(self, stmt: Select[Any]) -> list[MessageSearchResult]:
        async with self._session_factory() as conn:
            rows = (await conn.execute(stmt)).all()
        results: list[MessageSearchResult] = []
        for (
            message,
            session_title,
            session_host,
            session_tool,
            session_id,
            working_directory,
            rank,
        ) in rows:
            results.append(
                MessageSearchResult(
                    node_id=message.node_id,
                    parent_node_id=message.parent_node_id,
                    message_id=message.message_id,
                    role=message.role,
                    content=message.content,
                    created_at=message.created_at,
                    metadata=message.metadata_,
                    session_id=session_id,
                    session_title=session_title,
                    session_host=session_host,
                    session_tool=session_tool,
                    session_working_directory=working_directory,
                    rank=float(rank),
                )
            )
        return results

    async def _execute_session_search(self, stmt: Select[Any]) -> list[SessionRow]:
        async with self._session_factory() as conn:
            rows = (await conn.execute(stmt)).scalars().all()
        return [self._session_row(row) for row in rows]

    def _session_row(self, session: SessionModel) -> SessionRow:
        return SessionRow(
            host=session.host,
            tool=session.tool,
            session_id=session.session_id,
            title=session.title,
            working_directory=session.working_directory,
            model_name=session.model_name,
            created_at=session.created_at,
            last_activity_at=session.last_activity_at,
            message_count=session.message_count,
            hidden=session.hidden,
            synced_at=session.synced_at,
        )

    def _message_row(self, message: MessageModel) -> MessageRow:
        return MessageRow(
            node_id=message.node_id,
            parent_node_id=message.parent_node_id,
            message_id=message.message_id,
            role=message.role,
            content=message.content,
            created_at=message.created_at,
            metadata=message.metadata_,
        )
