from typing import Optional

from sqlalchemy import delete, func, select
from sqlalchemy.orm import Session

from vdash.config import Settings
from vdash.core import Database, NotFoundError, inject_session
from vdash.modules.user_db import UserDB
from vdash.pipelines.sql import IntentEnum, SqlAgent
from vdash.shared import PaginatedResult, PaginationQuery
from vdash.utils import CryptoHelper

from .enums import MessageRole
from .models import SqlConversation, SqlMessage, SqlQuery
from .schemas import (
    CreateSqlQueryRequest,
    GenerateSqlRequest,
    SqlContent,
    SqlMessageResponse,
    TextContent,
    UpdateSqlQueryRequest,
)


class SqlService:
    def __init__(
        self,
        db: Database,
        settings: Settings,
        sql_agent: SqlAgent,
        crypto_helper: CryptoHelper,
    ):
        self._db = db
        self._settings = settings
        self._sql_agent = sql_agent
        self._crypto_helper = crypto_helper

    # Conversation
    @inject_session
    async def list_conversations(
        self,
        pagination: PaginationQuery,
        user_db_uuid: Optional[str] = None,
        session: Optional[Session] = None,
    ) -> PaginatedResult[SqlConversation]:
        assert session is not None

        if pagination.page < 1 or pagination.page_size < 1:
            return PaginatedResult(
                items=[], page=pagination.page, page_size=pagination.page_size, total=0
            )

        offset = (pagination.page - 1) * pagination.page_size

        stmt = select(SqlConversation).join(
            UserDB, UserDB.id == SqlConversation.user_db_id
        )

        if user_db_uuid:
            stmt = stmt.where(UserDB.uuid == user_db_uuid)

        if pagination.q:
            stmt = stmt.where(SqlConversation.title.ilike(f"%{pagination.q}%"))

        count_stmt = select(func.count()).select_from(stmt.subquery())
        total = session.scalar(count_stmt) or 0

        stmt = (
            stmt.order_by(SqlConversation.created_at.desc())
            .offset(offset)
            .limit(pagination.page_size)
        )
        items = session.execute(stmt).scalars().all()

        return PaginatedResult(
            items=items,
            page=pagination.page,
            page_size=pagination.page_size,
            total=total,
        )

    @inject_session
    async def generate_sql(
        self,
        request: GenerateSqlRequest,
        session: Optional[Session] = None,
    ) -> SqlMessageResponse:
        """
        Generate SQL from a natural language question.

        Flow:
        1. Resolve user_db connection string and get or create a conversation.
        2. Route to SqlAgent: confirm if waiting for user, otherwise ask new question.
        3. Classify intent and build response content:
           - IRRELEVANT / AMBIGUOUS_SQL → TextContent
           - DIRECT_SQL → SqlContent
        4. Persist user and assistant messages, return response.
        """
        assert session is not None

        # Step 1: Resolve user_db and decrypt the stored connection string
        user_db = (
            session.execute(select(UserDB).where(UserDB.uuid == request.user_db_uuid))
            .scalars()
            .one_or_none()
        )
        if user_db is None:
            raise NotFoundError(f"Database not found: {request.user_db_uuid}")
        db_url = self._crypto_helper.decrypt(user_db.connection_string)

        # Step 1b: Get existing conversation or create a new one
        conversation = await self._get_or_create_conversation(
            request=request,
            user_db_id=user_db.id,
            session=session,
        )

        # Step 2: Route to SqlAgent — if the agent is waiting for user confirmation,
        # call confirm; otherwise start a new question
        if self._sql_agent.is_waiting_for_user(thread_id=conversation.uuid):
            agent_state = self._sql_agent.confirm_question(
                db_url=db_url,
                thread_id=conversation.uuid,
                user_input=request.user_input,
            )
        else:
            agent_state = self._sql_agent.ask_question(
                db_url=db_url,
                thread_id=conversation.uuid,
                user_input=request.user_input,
            )

        # Step 3: Classify intent and build the appropriate response content
        intent_class = agent_state.intent_classification
        if intent_class and intent_class.intent == IntentEnum.IRRELEVANT:
            # Irrelevant question → return a text message
            content = TextContent(text=intent_class.message_to_user or "")
        elif intent_class and intent_class.intent == IntentEnum.AMBIGUOUS_SQL:
            # Ambiguous question → return text with clarification suggestions
            content = TextContent(
                text=intent_class.message_to_user or "",
                suggestions=intent_class.clarification.suggestions
                if intent_class.clarification
                else None,
            )
        elif intent_class and intent_class.intent == IntentEnum.DIRECT_SQL:
            # Clear question → return the generated SQL
            content = SqlContent(sql=agent_state.query_sql or "")
        else:
            content = TextContent(
                text="I can't understand your question. Please try again."
            )

        # Step 4: Persist user and assistant messages to the conversation
        user_message = SqlMessage(
            role=MessageRole.USER,
            content=TextContent(text=request.user_input).model_dump(),
            conversation=conversation,
        )
        assistant_message = SqlMessage(
            role=MessageRole.ASSISTANT,
            content=content.model_dump(),
            conversation=conversation,
        )
        session.add_all([user_message, assistant_message])
        session.flush()

        return SqlMessageResponse(
            uuid=assistant_message.uuid,
            role=assistant_message.role,
            content=assistant_message.content,
            created_at=assistant_message.created_at,
        )

    @inject_session
    async def get_conversation(
        self,
        conversation_uuid: str,
        session: Optional[Session] = None,
    ) -> SqlConversation:
        """Get a single conversation by UUID."""
        assert session is not None

        conversation = (
            session.execute(
                select(SqlConversation).where(SqlConversation.uuid == conversation_uuid)
            )
            .scalars()
            .one_or_none()
        )

        if conversation is None:
            raise NotFoundError(f"Conversation not found: {conversation_uuid}")

        return conversation

    @inject_session
    async def delete_conversation(
        self,
        conversation_uuid: str,
        session: Optional[Session] = None,
    ) -> None:
        """Delete a conversation and all its messages."""
        assert session is not None

        conversation = (
            session.execute(
                select(SqlConversation).where(SqlConversation.uuid == conversation_uuid)
            )
            .scalars()
            .one_or_none()
        )

        if conversation is None:
            raise NotFoundError(f"Conversation not found: {conversation_uuid}")

        session.execute(
            delete(SqlMessage).where(SqlMessage.conversation_id == conversation.id)
        )
        session.delete(conversation)

    @inject_session
    async def list_messages(
        self,
        conversation_uuid: str,
        pagination: PaginationQuery,
        session: Optional[Session] = None,
    ) -> PaginatedResult[SqlMessage]:
        assert session is not None

        conversation = (
            session.execute(
                select(SqlConversation).where(SqlConversation.uuid == conversation_uuid)
            )
            .scalars()
            .one_or_none()
        )

        if conversation is None:
            raise NotFoundError(f"Conversation not found: {conversation_uuid}")

        if pagination.page < 1 or pagination.page_size < 1:
            return PaginatedResult(
                items=[], page=pagination.page, page_size=pagination.page_size, total=0
            )

        offset = (pagination.page - 1) * pagination.page_size

        stmt = select(SqlMessage).where(SqlMessage.conversation_id == conversation.id)

        count_stmt = select(func.count()).select_from(stmt.subquery())
        total = session.scalar(count_stmt) or 0

        stmt = (
            stmt.order_by(SqlMessage.created_at.desc())
            .offset(offset)
            .limit(pagination.page_size)
        )
        items = session.execute(stmt).scalars().all()

        return PaginatedResult(
            items=items,
            page=pagination.page,
            page_size=pagination.page_size,
            total=total,
        )

    # SqlQuery
    @inject_session
    async def list_queries(
        self,
        pagination: PaginationQuery,
        conversation_uuid: Optional[str] = None,
        session: Optional[Session] = None,
    ) -> PaginatedResult[SqlQuery]:
        """List saved SQL queries with conversation_uuid."""
        assert session is not None

        if pagination.page < 1 or pagination.page_size < 1:
            return PaginatedResult(
                items=[], page=pagination.page, page_size=pagination.page_size, total=0
            )

        offset = (pagination.page - 1) * pagination.page_size

        stmt = select(SqlQuery)

        if conversation_uuid:
            stmt = (
                stmt.join(SqlMessage, SqlMessage.id == SqlQuery.message_id)
                .join(
                    SqlConversation,
                    SqlConversation.id == SqlMessage.conversation_id,
                )
                .where(SqlConversation.uuid == conversation_uuid)
            )

        if pagination.q:
            stmt = stmt.where(SqlQuery.name.ilike(f"%{pagination.q}%"))

        count_stmt = select(func.count()).select_from(stmt.subquery())
        total = session.scalar(count_stmt) or 0

        stmt = (
            stmt.order_by(SqlQuery.created_at.desc())
            .offset(offset)
            .limit(pagination.page_size)
        )
        items = session.execute(stmt).scalars().all()

        return PaginatedResult(
            items=items,
            page=pagination.page,
            page_size=pagination.page_size,
            total=total,
        )

    @inject_session
    async def create_query(
        self,
        request: CreateSqlQueryRequest,
        session: Optional[Session] = None,
    ) -> SqlQuery:
        """Create a new SQL query."""
        assert session is not None

        message = (
            session.execute(
                select(SqlMessage).where(SqlMessage.uuid == request.message_uuid)
            )
            .scalars()
            .one_or_none()
        )

        if message is None:
            raise NotFoundError(f"Message not found: {request.message_uuid}")

        if message.content.get("type") != "sql":
            raise ValueError("Message does not contain SQL content")

        user_db = message.conversation.user_db

        query = SqlQuery(
            user_db_id=user_db.id,
            message_id=message.id,
            name=request.name,
            sql=request.sql,
        )
        session.add(query)
        session.flush()

        return query

    @inject_session
    async def get_query(
        self,
        query_uuid: str,
        session: Optional[Session] = None,
    ) -> SqlQuery:
        """Get a single SQL query by UUID."""
        assert session is not None

        query = (
            session.execute(select(SqlQuery).where(SqlQuery.uuid == query_uuid))
            .scalars()
            .one_or_none()
        )

        if query is None:
            raise NotFoundError(f"Query not found: {query_uuid}")

        return query

    @inject_session
    async def update_query(
        self,
        query_uuid: str,
        request: UpdateSqlQueryRequest,
        session: Optional[Session] = None,
    ) -> SqlQuery:
        """Update a SQL query's name, sql, and/or user_db."""
        assert session is not None

        query = (
            session.execute(select(SqlQuery).where(SqlQuery.uuid == query_uuid))
            .scalars()
            .one_or_none()
        )

        if query is None:
            raise NotFoundError(f"Query not found: {query_uuid}")

        if request.user_db_uuid is not None:
            user_db = (
                session.execute(
                    select(UserDB).where(UserDB.uuid == request.user_db_uuid)
                )
                .scalars()
                .one_or_none()
            )
            if user_db is None:
                raise NotFoundError(f"Database not found: {request.user_db_uuid}")
            query.user_db_id = user_db.id

        if request.name is not None:
            query.name = request.name
        if request.sql is not None:
            query.sql = request.sql
        session.flush()

        return query

    @inject_session
    async def delete_query(
        self,
        query_uuid: str,
        session: Optional[Session] = None,
    ) -> None:
        """Delete a SQL query by UUID."""
        assert session is not None

        query = (
            session.execute(select(SqlQuery).where(SqlQuery.uuid == query_uuid))
            .scalars()
            .one_or_none()
        )

        if query is None:
            raise NotFoundError(f"Query not found: {query_uuid}")

        session.delete(query)

    async def _get_or_create_conversation(
        self,
        request: GenerateSqlRequest,
        user_db_id: int,
        session: Session,
    ) -> SqlConversation:
        """Get existing conversation or create a new one."""
        if request.conversation_uuid:
            conversation = (
                session.execute(
                    select(SqlConversation).where(
                        SqlConversation.uuid == request.conversation_uuid
                    )
                )
                .scalars()
                .one_or_none()
            )

            if conversation is None:
                raise NotFoundError(
                    f"Conversation not found: {request.conversation_uuid}"
                )
            return conversation

        # Create new conversation
        conversation = SqlConversation(
            user_db_id=user_db_id, title=request.user_input[:500]
        )
        session.add(conversation)
        session.flush()
        return conversation
