import tiktoken
from loguru import logger
from openrouter import OpenRouter
from sqlalchemy import Float, func, literal, select, update
from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)

from session_hub_core.config import DatabaseConfig, EmbeddingConfig
from session_hub_core.exceptions import SessionNotFoundError, SyncError
from session_hub_core.models import (
    MessageModel,
    MessageSearchResult,
    SessionModel,
    SessionSearchResult,
)


def _truncate_text(text: str, max_tokens: int) -> str:
    try:
        encoding = tiktoken.get_encoding("cl100k_base")
        token_count = len(encoding.encode(text))
    except Exception:
        token_count = None
    if token_count is None or token_count <= max_tokens:
        return text
    return text[: max_tokens * 4]


class EmbeddingService:
    """Generate and store embeddings for sessions and messages.

    Uses OpenRouter for embedding generation (supports 400+ models from
    multiple providers — OpenAI, Qwen, Cohere, etc.). Uses pgvector for
    storage and HNSW indexes for similarity search. Embeddings are
    generated asynchronously after sync — they don't block the sync
    pipeline.
    """

    def __init__(self, db_config: DatabaseConfig, emb_config: EmbeddingConfig) -> None:
        self._db_config = db_config
        self._emb_config = emb_config
        self._engine: AsyncEngine = create_async_engine(
            db_config.dsn,
            pool_size=db_config.pool_size,
            max_overflow=db_config.max_overflow,
            pool_pre_ping=True,
            pool_recycle=300,
            echo=db_config.echo,
        )
        self._client = OpenRouter(api_key=emb_config.api_key, server_url=emb_config.base_url)
        self._session_factory: async_sessionmaker[AsyncSession] = async_sessionmaker(
            self._engine, expire_on_commit=False
        )

    async def embed_session(self, host: str, tool: str, session_id: str) -> None:
        """Generate and store the session-level embedding.

        Concatenates session title + all message content, truncates to
        max_tokens, calls OpenRouter, and stores the vector in
        sessions.embedding. Updates embedding_model and embedded_at.
        """
        async with self._session_factory() as conn:
            row = (
                await conn.execute(
                    select(SessionModel).where(
                        SessionModel.host == host,
                        SessionModel.tool == tool,
                        SessionModel.session_id == session_id,
                    )
                )
            ).scalar_one_or_none()
            if row is None:
                raise SessionNotFoundError(f"Session {host}/{tool}/{session_id} not found")
            contents = (
                (
                    await conn.execute(
                        select(MessageModel.content)
                        .where(MessageModel.session_pk == row.id)
                        .order_by(MessageModel.id)
                    )
                )
                .scalars()
                .all()
            )
            text_to_embed = row.title + "\n" + "\n".join(contents)
            vector = (await self.embed_batch([text_to_embed]))[0]
            await conn.execute(
                update(SessionModel)
                .where(SessionModel.id == row.id)
                .values(
                    embedding=vector,
                    embedding_model=self._emb_config.model,
                    embedded_at=func.now(),
                )
            )
            await conn.commit()

    async def embed_message(self, host: str, tool: str, session_id: str, node_id: str) -> None:
        """Generate and store a message-level embedding.

        Uses message content, calls OpenRouter, stores in messages.embedding.
        """
        async with self._session_factory() as conn:
            row = (
                await conn.execute(
                    select(MessageModel).where(
                        MessageModel.host == host,
                        MessageModel.tool == tool,
                        MessageModel.session_id == session_id,
                        MessageModel.node_id == node_id,
                    )
                )
            ).scalar_one_or_none()
            if row is None:
                raise SessionNotFoundError(
                    f"Message {host}/{tool}/{session_id}/{node_id} not found"
                )
            vector = (await self.embed_batch([row.content]))[0]
            await conn.execute(
                update(MessageModel)
                .where(MessageModel.id == row.id)
                .values(
                    embedding=vector,
                    embedding_model=self._emb_config.model,
                    embedded_at=func.now(),
                )
            )
            await conn.commit()

    async def embed_batch(self, texts: list[str]) -> list[list[float]]:
        """Call OpenRouter embeddings API for a batch of texts.

        Returns a list of vectors (dimension depends on the configured
        model — 1536 for openai/text-embedding-3-small).
        """
        truncated = [_truncate_text(t, self._emb_config.max_tokens_per_text) for t in texts]
        try:
            response = await self._client.embeddings.generate_async(
                input=truncated,
                model=self._emb_config.model,
                dimensions=self._emb_config.dimensions,
                encoding_format="float",
            )
        except Exception as exc:
            raise SyncError(f"Embedding API call failed: {exc}") from exc
        if isinstance(response, str):
            raise SyncError(f"Unexpected response from embeddings API: {response[:200]}")
        vectors: list[list[float]] = []
        for item in response.data:
            if isinstance(item.embedding, str):
                raise SyncError(
                    "Embedding API returned base64 encoding despite encoding_format='float'"
                )
            vectors.append(item.embedding)
        return vectors

    async def backfill_missing(
        self,
        batch_size: int = 100,
        only_sessions: bool = False,
        only_messages: bool = False,
    ) -> int:
        """Backfill embeddings for rows where embedding IS NULL.

        Called by the pg_cron job or manually. Returns count of rows
        embedded.
        """
        count = 0
        if not only_messages:
            count += await self._backfill_sessions(batch_size)
        if not only_sessions:
            count += await self._backfill_messages(batch_size)
        return count

    async def _backfill_sessions(self, batch_size: int) -> int:
        embedded = 0
        while True:
            async with self._session_factory() as conn:
                rows = (
                    (
                        await conn.execute(
                            select(SessionModel)
                            .where(SessionModel.embedding.is_(None))
                            .order_by(SessionModel.id)
                            .limit(batch_size)
                        )
                    )
                    .scalars()
                    .all()
                )
                if not rows:
                    return embedded
                for session in rows:
                    contents = (
                        (
                            await conn.execute(
                                select(MessageModel.content)
                                .where(MessageModel.session_pk == session.id)
                                .order_by(MessageModel.id)
                            )
                        )
                        .scalars()
                        .all()
                    )
                    text_to_embed = session.title + "\n" + "\n".join(contents)
                    vector = (await self.embed_batch([text_to_embed]))[0]
                    await conn.execute(
                        update(SessionModel)
                        .where(SessionModel.id == session.id)
                        .values(
                            embedding=vector,
                            embedding_model=self._emb_config.model,
                            embedded_at=func.now(),
                        )
                    )
                    embedded += 1
                await conn.commit()
            logger.info("Backfilled {} session embeddings", embedded)

    async def _backfill_messages(self, batch_size: int) -> int:
        embedded = 0
        while True:
            async with self._session_factory() as conn:
                rows = (
                    (
                        await conn.execute(
                            select(MessageModel)
                            .where(MessageModel.embedding.is_(None))
                            .order_by(MessageModel.id)
                            .limit(batch_size)
                        )
                    )
                    .scalars()
                    .all()
                )
                if not rows:
                    return embedded
                vectors = await self.embed_batch([m.content for m in rows])
                for message, vector in zip(rows, vectors, strict=True):
                    await conn.execute(
                        update(MessageModel)
                        .where(MessageModel.id == message.id)
                        .values(
                            embedding=vector,
                            embedding_model=self._emb_config.model,
                            embedded_at=func.now(),
                        )
                    )
                    embedded += 1
                await conn.commit()
            logger.info("Backfilled {} message embeddings", embedded)

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

        Uses HNSW index with cosine distance:
        SELECT ... FROM sessions
        WHERE embedding IS NOT NULL
        ORDER BY embedding <=> $query_vector
        LIMIT $limit
        """
        vector = (await self.embed_batch([query]))[0]
        distance = SessionModel.embedding.op("<=>")(vector)
        stmt = (
            select(SessionModel, (literal(1.0, Float) - distance).label("similarity"))
            .where(SessionModel.embedding.is_not(None))
            .order_by(distance)
            .limit(limit)
        )
        if tool is not None:
            stmt = stmt.where(SessionModel.tool == tool)
        if host is not None:
            stmt = stmt.where(SessionModel.host == host)
        async with self._session_factory() as conn:
            rows = (await conn.execute(stmt)).all()
        results: list[SessionSearchResult] = []
        for session, similarity in rows:
            results.append(
                SessionSearchResult(
                    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,
                    similarity=float(similarity),
                    matched_text=session.title[:200],
                )
            )
        return results

    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: embed the query, find nearest messages.

        Uses HNSW index with cosine distance on messages.embedding.
        """
        vector = (await self.embed_batch([query]))[0]
        distance = MessageModel.embedding.op("<=>")(vector)
        stmt = (
            select(
                MessageModel,
                SessionModel.title,
                SessionModel.session_id,
                SessionModel.working_directory,
                (literal(1.0, Float) - distance).label("similarity"),
            )
            .join(SessionModel, SessionModel.id == MessageModel.session_pk)
            .where(MessageModel.embedding.is_not(None))
            .order_by(distance)
            .limit(limit)
        )
        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)
        async with self._session_factory() as conn:
            rows = (await conn.execute(stmt)).all()
        results: list[MessageSearchResult] = []
        for message, session_title, session_id, session_working_directory, similarity 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=message.host,
                    session_tool=message.tool,
                    session_working_directory=session_working_directory,
                    rank=float(similarity),
                )
            )
        return results

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