from collections.abc import Sequence
from logging import getLogger
from typing import List, Optional

from dotenv import load_dotenv
from openai import AsyncOpenAI
from sqlalchemy import Select, cast, insert, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.sql import func
from sqlalchemy.types import BigInteger

from . import models, schemas
from .exceptions import (
    ConflictException,
    ResourceNotFoundException,
    ValidationException,
)

load_dotenv(override=True)

openai_client = AsyncOpenAI()

logger = getLogger(__name__)

DEF_PROTECTED_COLLECTION_NAME = "honcho"

########################################################
# app methods
########################################################


async def get_app(db: AsyncSession, app_id: str) -> models.App:
    """
    Get an app by its ID.

    Args:
        db: Database session
        app_id: Public ID of the app

    Returns:
        The app if found

    Raises:
        ResourceNotFoundException: If the app does not exist
    """
    stmt = select(models.App).where(models.App.public_id == app_id)
    result = await db.execute(stmt)
    app = result.scalar_one_or_none()
    if app is None:
        logger.warning(f"App with ID {app_id} not found")
        raise ResourceNotFoundException(f"App with ID {app_id} not found")
    return app


async def get_all_apps(
    db: AsyncSession,
    reverse: Optional[bool] = False,
    filter: Optional[dict] = None,
) -> Select:
    """
    Get all apps.

    Args:
        db: Database session
        reverse: Whether to reverse the order of the apps
        filter: Filter the apps by a dictionary of metadata
    """
    stmt = select(models.App)
    if reverse:
        stmt = stmt.order_by(models.App.id.desc())
    else:
        stmt = stmt.order_by(models.App.id)
    if filter is not None:
        stmt = stmt.where(models.App.h_metadata.contains(filter))
    return stmt


async def get_app_by_name(db: AsyncSession, name: str) -> models.App:
    """
    Get an app by its name.

    Args:
        db: Database session
        name: Name of the app

    Returns:
        The app if found

    Raises:
        ResourceNotFoundException: If the app does not exist
    """
    stmt = select(models.App).where(models.App.name == name)
    result = await db.execute(stmt)
    app = result.scalar_one_or_none()
    if app is None:
        logger.warning(f"App with name '{name}' not found")
        raise ResourceNotFoundException(f"App with name '{name}' not found")
    return app


# def get_apps(db: AsyncSession) -> Sequence[models.App]:
#     return db.query(models.App).all()


async def create_app(db: AsyncSession, app: schemas.AppCreate) -> models.App:
    """
    Create a new app.

    Args:
        db: Database session
        app: App creation schema

    Returns:
        The created app

    Raises:
        ConflictException: If an app with the same name already exists
    """
    try:
        honcho_app = models.App(name=app.name, h_metadata=app.metadata)
        db.add(honcho_app)
        await db.commit()
        logger.info(f"App created successfully: {app.name}")
        return honcho_app
    except IntegrityError as e:
        await db.rollback()
        logger.error(f"IntegrityError creating app with name '{app.name}': {str(e)}")
        raise ConflictException(f"App with name '{app.name}' already exists") from e


async def update_app(
    db: AsyncSession, app_id: str, app: schemas.AppUpdate
) -> models.App:
    """
    Update an app.

    Args:
        db: Database session
        app_id: Public ID of the app
        app: App update schema

    Returns:
        The updated app

    Raises:
        ResourceNotFoundException: If the app does not exist
    """
    try:
        honcho_app = await get_app(db, app_id)

        if app.name is not None:
            honcho_app.name = app.name
        if app.metadata is not None:
            honcho_app.h_metadata = app.metadata

        await db.commit()
        logger.info(f"App with ID {app_id} updated successfully")
        return honcho_app
    except IntegrityError as e:
        await db.rollback()
        logger.error(f"IntegrityError updating app {app_id}: {str(e)}")
        raise ConflictException(
            "App update failed - unique constraint violation"
        ) from e


########################################################
# user methods
########################################################


async def create_user(
    db: AsyncSession, app_id: str, user: schemas.UserCreate
) -> models.User:
    """
    Create a new user.

    Args:
        db: Database session
        app_id: Public ID of the app
        user: User creation schema

    Returns:
        The created user

    Raises:
        ConflictException: If a user with the same name already exists in this app
    """
    try:
        honcho_user = models.User(
            app_id=app_id,
            name=user.name,
            h_metadata=user.metadata,
        )
        db.add(honcho_user)
        await db.commit()
        logger.info(f"User created successfully: {user.name} for app {app_id}")
        return honcho_user
    except IntegrityError as e:
        await db.rollback()
        logger.warning(f"Failed to create user - integrity error: {str(e)}")
        raise ConflictException("User with this name already exists") from e


async def get_user(db: AsyncSession, app_id: str, user_id: str) -> models.User:
    """
    Get a user by app ID and user ID.

    Args:
        db: Database session
        app_id: Public ID of the app
        user_id: Public ID of the user

    Returns:
        The user if found

    Raises:
        ResourceNotFoundException: If the user does not exist
    """
    stmt = (
        select(models.User)
        .where(models.User.app_id == app_id)
        .where(models.User.public_id == user_id)
    )
    result = await db.execute(stmt)
    user = result.scalar_one_or_none()
    if user is None:
        logger.warning(f"User with ID '{user_id}' not found in app {app_id}")
        raise ResourceNotFoundException(f"User with ID '{user_id}' not found")
    return user


async def get_user_by_name(db: AsyncSession, app_id: str, name: str) -> models.User:
    """
    Get a user by app ID and name.

    Args:
        db: Database session
        app_id: Public ID of the app
        name: Name of the user

    Returns:
        The user if found

    Raises:
        ResourceNotFoundException: If the user does not exist
    """
    stmt = (
        select(models.User)
        .where(models.User.app_id == app_id)
        .where(models.User.name == name)
    )
    result = await db.execute(stmt)
    user = result.scalar_one_or_none()
    if user is None:
        logger.warning(f"User with name '{name}' not found in app {app_id}")
        raise ResourceNotFoundException(f"User with name '{name}' not found")
    return user


async def get_users(
    db: AsyncSession,
    app_id: str,
    reverse: bool = False,
    filter: Optional[dict] = None,
) -> Select:
    stmt = select(models.User).where(models.User.app_id == app_id)

    if filter is not None:
        stmt = stmt.where(models.User.h_metadata.contains(filter))

    if reverse:
        stmt = stmt.order_by(models.User.id.desc())
    else:
        stmt = stmt.order_by(models.User.id)

    return stmt


async def update_user(
    db: AsyncSession, app_id: str, user_id: str, user: schemas.UserUpdate
) -> models.User:
    """
    Update a user.

    Args:
        db: Database session
        app_id: Public ID of the app
        user_id: Public ID of the user
        user: User update schema

    Returns:
        The updated user

    Raises:
        ResourceNotFoundException: If the user does not exist
        ValidationException: If the update data is invalid
        ConflictException: If the update violates a unique constraint
    """
    try:
        # get_user will raise ResourceNotFoundException if not found
        honcho_user = await get_user(db, app_id, user_id)

        if user.name is not None:
            honcho_user.name = user.name
        if user.metadata is not None:
            honcho_user.h_metadata = user.metadata

        await db.commit()
        logger.info(f"User {user_id} updated successfully")
        return honcho_user
    except IntegrityError as e:
        await db.rollback()
        logger.warning(f"User update failed due to integrity error: {str(e)}")
        raise ConflictException(
            "User update failed - unique constraint violation"
        ) from e


########################################################
# session methods
########################################################


async def get_session(
    db: AsyncSession,
    app_id: str,
    session_id: str,
    user_id: Optional[str] = None,
) -> models.Session:
    """
    Get a session by ID for a specific user and app.

    Args:
        db: Database session
        app_id: Public ID of the app
        session_id: Public ID of the session
        user_id: Optional public ID of the user

    Returns:
        The session if found

    Raises:
        ResourceNotFoundException: If the session does not exist or doesn't belong to the user
    """
    stmt = (
        select(models.Session)
        .where(models.Session.app_id == app_id)
        .where(models.Session.public_id == session_id)
    )
    if user_id is not None:
        stmt = stmt.where(models.Session.user_id == user_id)
    result = await db.execute(stmt)
    session = result.scalar_one_or_none()
    if session is None:
        logger.warning(f"Session with ID '{session_id}' not found for user {user_id}")
        raise ResourceNotFoundException("Session not found or does not belong to user")
    return session


async def get_sessions(
    app_id: str,
    user_id: str,
    reverse: Optional[bool] = False,
    is_active: Optional[bool] = False,
    filter: Optional[dict] = None,
) -> Select:
    stmt = (
        select(models.Session)
        .where(models.Session.app_id == app_id)
        .where(models.Session.user_id == user_id)
    )

    if is_active:
        stmt = stmt.where(models.Session.is_active.is_(True))

    if filter is not None:
        stmt = stmt.where(models.Session.h_metadata.contains(filter))

    if reverse:
        stmt = stmt.order_by(models.Session.id.desc())
    else:
        stmt = stmt.order_by(models.Session.id)

    return stmt


async def create_session(
    db: AsyncSession,
    session: schemas.SessionCreate,
    app_id: str,
    user_id: str,
) -> models.Session:
    """
    Create a new session for a user.

    Args:
        db: Database session
        session: Session creation schema
        app_id: ID of the app
        user_id: ID of the user

    Returns:
        The created session

    Raises:
        ResourceNotFoundException: If the user does not exist
    """
    try:
        # This will raise ResourceNotFoundException if user not found
        _honcho_user = await get_user(db, app_id=app_id, user_id=user_id)

        honcho_session = models.Session(
            user_id=user_id,
            app_id=app_id,
            h_metadata=session.metadata,
        )
        db.add(honcho_session)
        await db.commit()
        logger.info(f"Session created successfully for user {user_id}")
        return honcho_session
    except Exception as e:
        await db.rollback()
        logger.error(f"Error creating session for user {user_id}: {str(e)}")
        raise


async def update_session(
    db: AsyncSession,
    session: schemas.SessionUpdate,
    app_id: str,
    user_id: str,
    session_id: str,
) -> models.Session:
    """
    Update a session.

    Args:
        db: Database session
        session: Session update schema
        app_id: ID of the app
        user_id: ID of the user
        session_id: ID of the session

    Returns:
        The updated session

    Raises:
        ResourceNotFoundException: If the session does not exist or doesn't belong to the user
    """
    honcho_session = await get_session(
        db, app_id=app_id, session_id=session_id, user_id=user_id
    )
    if honcho_session is None:
        logger.warning(f"Session {session_id} not found for user {user_id}")
        raise ResourceNotFoundException("Session not found or does not belong to user")

    if (
        session.metadata is not None
    ):  # Need to explicitly be there won't make it empty by default
        honcho_session.h_metadata = session.metadata

    await db.commit()
    logger.info(f"Session {session_id} updated successfully")
    return honcho_session


async def delete_session(
    db: AsyncSession, app_id: str, user_id: str, session_id: str
) -> bool:
    """
    Mark a session as inactive (soft delete).

    Args:
        db: Database session
        app_id: ID of the app
        user_id: ID of the user
        session_id: ID of the session

    Returns:
        True if the session was deleted successfully

    Raises:
        ResourceNotFoundException: If the session does not exist or doesn't belong to the user
    """
    stmt = (
        select(models.Session)
        .where(models.Session.public_id == session_id)
        .where(models.Session.user_id == user_id)
        .where(models.Session.app_id == app_id)
    )
    result = await db.execute(stmt)
    honcho_session = result.scalar_one_or_none()

    if honcho_session is None:
        logger.warning(f"Session {session_id} not found for user {user_id}")
        raise ResourceNotFoundException("Session not found or does not belong to user")

    honcho_session.is_active = False
    await db.commit()
    logger.info(f"Session {session_id} marked as inactive")
    return True


async def clone_session(
    db: AsyncSession,
    app_id: str,
    user_id: str,
    original_session_id: str,
    cutoff_message_id: Optional[str] = None,
    deep_copy: bool = True,
) -> models.Session:
    """
    Clone a session and its messages. If cutoff_message_id is provided,
    only clone messages up to and including that message.

    Args:
        db: SQLAlchemy session
        app_id: ID of the app the target session is in
        user_id: ID of the user the target session belongs to
        original_session_id: ID of the session to clone
        cutoff_message_id: Optional ID of the last message to include in the clone

    Returns:
        The newly created session
    """
    # Get the original session
    stmt = (
        select(models.Session)
        .where(models.Session.app_id == app_id)
        .where(models.Session.user_id == user_id)
        .where(models.Session.public_id == original_session_id)
    )
    original_session = await db.scalar(stmt)
    if not original_session:
        raise ValueError("Original session not found")

    # If cutoff_message_id is provided, verify it belongs to the session
    cutoff_message = None
    if cutoff_message_id is not None:
        stmt = select(models.Message).where(
            models.Message.public_id == cutoff_message_id,
            models.Message.session_id == original_session_id,
        )
        cutoff_message = await db.scalar(stmt)
        if not cutoff_message:
            raise ValueError(
                "Message not found or doesn't belong to the specified session"
            )

    # Create new session
    new_session = models.Session(
        user_id=original_session.user_id,
        app_id=original_session.app_id,
        h_metadata=original_session.h_metadata,
    )
    db.add(new_session)
    await db.flush()  # Flush to get the new session ID

    # Build query for messages to clone
    stmt = select(models.Message).where(
        models.Message.session_id == original_session_id
    )
    if cutoff_message_id is not None and cutoff_message is not None:
        stmt = stmt.where(models.Message.id <= cast(cutoff_message.id, BigInteger))
    stmt = stmt.order_by(models.Message.id)

    # Fetch messages to clone
    messages_to_clone_scalars = await db.scalars(stmt)
    messages_to_clone = messages_to_clone_scalars.all()

    if not messages_to_clone:
        return new_session

    # Prepare bulk insert data
    new_messages = [
        {
            "session_id": new_session.public_id,
            "content": message.content,
            "is_user": message.is_user,
            "h_metadata": message.h_metadata,
            "app_id": original_session.app_id,
            "user_id": original_session.user_id,
        }
        for message in messages_to_clone
    ]

    stmt = insert(models.Message).returning(models.Message.public_id)
    result = await db.execute(stmt, new_messages)
    new_message_ids = result.scalars().all()

    # Create mapping of old to new message IDs
    message_id_map = dict(
        zip([message.public_id for message in messages_to_clone], new_message_ids)
    )

    # Handle metamessages if deep copy is requested
    if deep_copy:
        # Fetch all metamessages tied to the session in a single query
        stmt = (
            select(models.Metamessage)
            .where(models.Metamessage.session_id == original_session_id)
            .order_by(models.Metamessage.id)  # Explicit ordering by id
        )
        if cutoff_message_id is not None and cutoff_message is not None:
            # Only get metamessages related to messages we're cloning
            message_ids = [message.public_id for message in messages_to_clone]
            stmt = stmt.where(
                (models.Metamessage.message_id.is_(None))
                | (models.Metamessage.message_id.in_(message_ids))
            )

        metamessages_result = await db.scalars(stmt)
        metamessages = metamessages_result.all()

        if metamessages:
            # Prepare bulk insert data for metamessages
            new_metamessages = []

            for meta in metamessages:
                # Base metamessage data
                meta_data = {
                    "user_id": meta.user_id,  # Preserve original user
                    "session_id": new_session.public_id,
                    "metamessage_type": meta.metamessage_type,
                    "content": meta.content,
                    "h_metadata": meta.h_metadata,
                    "app_id": original_session.app_id,
                }

                # If the metamessage was tied to a message, tie it to the corresponding new message
                if meta.message_id is not None and meta.message_id in message_id_map:
                    meta_data["message_id"] = message_id_map[meta.message_id]

                new_metamessages.append(meta_data)

            # Bulk insert metamessages using modern insert syntax
            if new_metamessages:
                stmt = insert(models.Metamessage)
                await db.execute(stmt, new_metamessages)

    await db.commit()

    return new_session


########################################################
# Message Methods
########################################################


async def create_message(
    db: AsyncSession,
    message: schemas.MessageCreate,
    app_id: str,
    user_id: str,
    session_id: str,
) -> models.Message:
    honcho_session = await get_session(
        db, app_id=app_id, session_id=session_id, user_id=user_id
    )
    if honcho_session is None:
        raise ValueError("Session not found or does not belong to user")

    honcho_message = models.Message(
        session_id=session_id,
        is_user=message.is_user,
        content=message.content,
        h_metadata=message.metadata,
        user_id=user_id,
        app_id=app_id,
    )
    db.add(honcho_message)
    await db.commit()
    # await db.refresh(honcho_message, attribute_names=["id", "content", "h_metadata"])
    # await db.refresh(honcho_message)
    return honcho_message


async def create_messages(
    db: AsyncSession,
    messages: list[schemas.MessageCreate],
    app_id: str,
    user_id: str,
    session_id: str,
) -> list[models.Message]:
    """Bulk create messages for a session while maintaining order"""
    # Verify session exists and belongs to user
    honcho_session = await get_session(
        db, app_id=app_id, session_id=session_id, user_id=user_id
    )
    if honcho_session is None:
        raise ValueError("Session not found or does not belong to user")

    # Create list of message records
    message_records = [
        {
            "session_id": session_id,
            "is_user": message.is_user,
            "content": message.content,
            "h_metadata": message.metadata,
            "user_id": user_id,
            "app_id": app_id,
        }
        for message in messages
    ]

    # Bulk insert messages and return them in order
    stmt = insert(models.Message).returning(models.Message)
    result = await db.execute(stmt, message_records)
    await db.commit()

    return list(result.scalars().all())


async def get_messages(
    db: AsyncSession,
    app_id: str,
    user_id: str,
    session_id: str,
    reverse: Optional[bool] = False,
    filter: Optional[dict] = None,
) -> Select:
    stmt = (
        select(models.Message)
        .where(models.Message.app_id == app_id)
        .where(models.Message.user_id == user_id)
        .where(models.Message.session_id == session_id)
    )

    if filter is not None:
        stmt = stmt.where(models.Message.h_metadata.contains(filter))

    if reverse:
        stmt = stmt.order_by(models.Message.id.desc())
    else:
        stmt = stmt.order_by(models.Message.id)

    return stmt


async def get_message(
    db: AsyncSession,
    app_id: str,
    user_id: str,
    session_id: str,
    message_id: str,
) -> Optional[models.Message]:
    stmt = (
        select(models.Message)
        .where(models.Message.app_id == app_id) 
        .where(models.Message.user_id == user_id)
        .where(models.Message.session_id == session_id)
        .where(models.Message.public_id == message_id)
    )
    result = await db.execute(stmt)
    return result.scalar_one_or_none()


async def update_message(
    db: AsyncSession,
    message: schemas.MessageUpdate,
    app_id: str,
    user_id: str,
    session_id: str,
    message_id: str,
) -> bool:
    honcho_message = await get_message(
        db, app_id=app_id, session_id=session_id, user_id=user_id, message_id=message_id
    )
    if honcho_message is None:
        raise ValueError("Message not found or does not belong to user")
    if (
        message.metadata is not None
    ):  # Need to explicitly be there won't make it empty by default
        honcho_message.h_metadata = message.metadata
    await db.commit()
    # await db.refresh(honcho_message)
    return honcho_message


########################################################
# metamessage methods
########################################################


async def create_metamessage(
    db: AsyncSession,
    user_id: str,
    metamessage: schemas.MetamessageCreate,
    app_id: str,
):
    # Validate user exists
    user = await get_user(db, app_id=app_id, user_id=user_id)
    if user is None:
        raise ResourceNotFoundException(f"User with ID '{user_id}' not found")

    # Initialize metamessage data
    metamessage_data = {
        "user_id": user_id,
        "app_id": app_id,
        "metamessage_type": metamessage.metamessage_type,
        "content": metamessage.content,
        "h_metadata": metamessage.metadata,
    }

    # Validate session_id if provided
    if metamessage.session_id is not None:
        session = await get_session(
            db,
            app_id=app_id,
            user_id=user_id,
            session_id=metamessage.session_id,
        )
        if session is None:
            raise ResourceNotFoundException(
                "Session not found or does not belong to user"
            )
        metamessage_data["session_id"] = metamessage.session_id

        # Validate message_id if provided
        if metamessage.message_id is not None:
            message = await get_message(
                db,
                app_id=app_id,
                session_id=metamessage.session_id,
                user_id=user_id,
                message_id=metamessage.message_id,
            )
            if message is None:
                raise ResourceNotFoundException(
                    "Message not found or does not belong to session"
                )
            metamessage_data["message_id"] = metamessage.message_id
    elif metamessage.message_id is not None:
        # If message_id provided but no session_id, that's an error
        raise ValidationException("Cannot specify message_id without session_id")

    # Create metamessage
    honcho_metamessage = models.Metamessage(**metamessage_data)
    db.add(honcho_metamessage)
    await db.commit()
    return honcho_metamessage


async def get_metamessages(
    db: AsyncSession,
    app_id: str,
    user_id: str,
    session_id: Optional[str] = None,
    message_id: Optional[str] = None,
    metamessage_type: Optional[str] = None,
    filter: Optional[dict] = None,
    reverse: Optional[bool] = False,
) -> Select:
    # Base query starts with metamessage and user relationship
    stmt = (
        select(models.Metamessage)
        .where(models.Metamessage.app_id == app_id)
        .where(models.Metamessage.user_id == user_id)
    )

    # If session_id is provided, filter by it
    if session_id is not None:
        stmt = stmt.where(models.Metamessage.session_id == session_id)

    # If message_id is provided, filter by it
    if message_id is not None:
        stmt = stmt.where(models.Metamessage.message_id == message_id)

    # Filter by metamessage_type if provided
    if metamessage_type is not None:
        stmt = stmt.where(models.Metamessage.metamessage_type == metamessage_type)

    # Apply metadata filter if provided
    if filter is not None:
        stmt = stmt.where(models.Metamessage.h_metadata.contains(filter))

    # Apply sort order
    if reverse:
        stmt = stmt.order_by(models.Metamessage.id.desc())
    else:
        stmt = stmt.order_by(models.Metamessage.id)

    return stmt


async def get_metamessage(
    db: AsyncSession,
    app_id: str,
    user_id: str,
    metamessage_id: str,
    session_id: Optional[str] = None,
    message_id: Optional[str] = None,
) -> Optional[models.Metamessage]:
    # Base query for metamessage by ID
    stmt = (
        select(models.Metamessage)
        .where(models.Metamessage.app_id == app_id)
        .where(models.Metamessage.user_id == user_id)
        .where(models.Metamessage.public_id == metamessage_id)
    )

    # Add optional filters
    if session_id is not None:
        stmt = stmt.where(models.Metamessage.session_id == session_id)
    if message_id is not None:
        stmt = stmt.where(models.Metamessage.message_id == message_id)

    result = await db.execute(stmt)
    return result.scalar_one_or_none()


async def update_metamessage(
    db: AsyncSession,
    metamessage: schemas.MetamessageUpdate,
    app_id: str,
    user_id: str,
    metamessage_id: str,
) -> bool:
    # First retrieve the metamessage
    metamessage_obj = await get_metamessage(
        db, app_id=app_id, user_id=user_id, metamessage_id=metamessage_id
    )
    if metamessage_obj is None:
        raise ResourceNotFoundException(
            f"Metamessage with ID {metamessage_id} not found"
        )

    # Validate the consistency of relationships if they're being changed
    # If we're setting message_id, we must have a session_id
    if metamessage.message_id is not None and metamessage.session_id is None:
        # If updating message_id but not session_id, use the existing session_id
        metamessage.session_id = metamessage_obj.session_id
        if metamessage.session_id is None:
            raise ValidationException("Cannot specify message_id without session_id")

    # If we're updating session_id and message_id, validate they belong together
    if metamessage.session_id is not None and metamessage.message_id is not None:
        message = await get_message(
            db,
            app_id=app_id,
            session_id=metamessage.session_id,
            user_id=metamessage.user_id,
            message_id=metamessage.message_id,
        )
        if message is None:
            raise ResourceNotFoundException(
                "Message not found or doesn't belong to session"
            )

    # Update fields
    if metamessage.session_id is not None:
        metamessage_obj.session_id = metamessage.session_id

    if metamessage.message_id is not None:
        metamessage_obj.message_id = metamessage.message_id

    if metamessage.metadata is not None:
        metamessage_obj.h_metadata = metamessage.metadata

    if metamessage.metamessage_type is not None:
        metamessage_obj.metamessage_type = metamessage.metamessage_type

    await db.commit()
    return metamessage_obj


########################################################
# collection methods
########################################################

# Should be very similar to the session methods


async def get_collections(
    db: AsyncSession,
    app_id: str,
    user_id: str,
    reverse: Optional[bool] = False,
    filter: Optional[dict] = None,
) -> Select:
    """Get a distinct list of the names of collections associated with a user"""
    stmt = (
        select(models.Collection)
        .where(models.Collection.app_id == app_id)
        .where(models.Collection.user_id == user_id)
    )

    if filter is not None:
        stmt = stmt.where(models.Collection.h_metadata.contains(filter))

    if reverse:
        stmt = stmt.order_by(models.Collection.id.desc())
    else:
        stmt = stmt.order_by(models.Collection.id)

    return stmt


async def get_collection_by_id(
    db: AsyncSession, app_id: str, user_id: str, collection_id: str
) -> models.Collection:
    """
    Get a collection by ID for a specific user and app.

    Args:
        db: Database session
        app_id: Public ID of the app
        user_id: Public ID of the user
        collection_id: Public ID of the collection

    Returns:
        The collection if found

    Raises:
        ResourceNotFoundException: If the collection does not exist
    """
    stmt = (
        select(models.Collection)
        .where(models.Collection.app_id == app_id)
        .where(models.Collection.user_id == user_id)
        .where(models.Collection.public_id == collection_id)
    )
    result = await db.execute(stmt)
    collection = result.scalar_one_or_none()
    if collection is None:
        logger.warning(
            f"Collection with ID '{collection_id}' not found for user {user_id}"
        )
        raise ResourceNotFoundException(
            "Collection not found or does not belong to user"
        )
    return collection


async def get_collection_by_name(
    db: AsyncSession, app_id: str, user_id: str, name: str
) -> models.Collection:
    """
    Get a collection by name for a specific user and app.

    Args:
        db: Database session
        app_id: Public ID of the app
        user_id: Public ID of the user
        name: Name of the collection

    Returns:
        The collection if found

    Raises:
        ResourceNotFoundException: If the collection does not exist
    """
    stmt = (
        select(models.Collection)
        .where(models.Collection.app_id == app_id)
        .where(models.Collection.user_id == user_id)
        .where(models.Collection.name == name)
    )
    result = await db.execute(stmt)
    collection = result.scalar_one_or_none()
    if collection is None:
        logger.warning(f"Collection with name '{name}' not found for user {user_id}")
        raise ResourceNotFoundException(f"Collection with name '{name}' not found")
    return collection


async def create_collection(
    db: AsyncSession,
    collection: schemas.CollectionCreate,
    app_id: str,
    user_id: str,
) -> models.Collection:
    """
    Create a new collection for a user.

    Args:
        db: Database session
        collection: Collection creation schema
        app_id: ID of the app
        user_id: ID of the user

    Returns:
        The created collection

    Raises:
        ConflictException: If a collection with the same name already exists for this user
        ValidationException: If the collection configuration is invalid
        ResourceNotFoundException: If the user does not exist
    """
    try:
        # This will raise ResourceNotFoundException if user not found
        await get_user(db, app_id=app_id, user_id=user_id)

        # Check for reserved names
        if collection.name == "honcho":
            logger.warning(
                f"Attempted to create collection with reserved name 'honcho' for user {user_id}"
            )
            raise ValidationException(
                "Invalid collection configuration - 'honcho' is a reserved name"
            )

        honcho_collection = models.Collection(
            user_id=user_id,
            app_id=app_id,
            name=collection.name,
            h_metadata=collection.metadata,
        )
        db.add(honcho_collection)
        await db.commit()
        logger.info(
            f"Collection '{collection.name}' created successfully for user {user_id}"
        )
        return honcho_collection
    except IntegrityError as e:
        await db.rollback()
        logger.warning(f"Failed to create collection - integrity error: {str(e)}")
        raise ConflictException(
            f"Collection with name '{collection.name}' already exists"
        ) from e


async def create_user_protected_collection(
    db: AsyncSession,
    app_id: str,
    user_id: str,
) -> models.Collection:
    honcho_collection = models.Collection(
        user_id=user_id,
        app_id=app_id,
        name=DEF_PROTECTED_COLLECTION_NAME,
    )
    try:
        db.add(honcho_collection)
        await db.commit()
    except IntegrityError:
        await db.rollback()
        raise ValueError("Collection already exists") from None
    return honcho_collection


async def get_or_create_user_protected_collection(
    db: AsyncSession,
    app_id: str,
    user_id: str,
) -> models.Collection:
    try:
        honcho_collection = await get_collection_by_name(
            db, app_id, user_id, DEF_PROTECTED_COLLECTION_NAME
        )
        return honcho_collection
    except ResourceNotFoundException:
        honcho_collection = await create_user_protected_collection(db, app_id, user_id)
        return honcho_collection


async def update_collection(
    db: AsyncSession,
    collection: schemas.CollectionUpdate,
    app_id: str,
    user_id: str,
    collection_id: str,
) -> models.Collection:
    """
    Update a collection.

    Args:
        db: Database session
        collection: Collection update schema
        app_id: ID of the app
        user_id: ID of the user
        collection_id: ID of the collection

    Returns:
        The updated collection

    Raises:
        ResourceNotFoundException: If the collection does not exist
        ValidationException: If the update data is invalid
        ConflictException: If the update violates a unique constraint
    """
    try:
        # Validate input
        if collection.name is None and collection.metadata is None:
            logger.warning(
                f"Collection update attempted with no fields provided for collection {collection_id}"
            )
            raise ValidationException(
                "Invalid collection configuration - at least one field must be provided"
            )

        # This will raise ResourceNotFoundException if not found
        honcho_collection = await get_collection_by_id(
            db, app_id=app_id, user_id=user_id, collection_id=collection_id
        )

        # Check for reserved names if name is being updated
        if collection.name == "honcho":
            logger.warning(
                f"Attempted to rename collection to reserved name 'honcho' for user {user_id}"
            )
            raise ValidationException(
                "Invalid collection configuration - 'honcho' is a reserved name"
            )

        if collection.metadata is not None:
            honcho_collection.h_metadata = collection.metadata

        if collection.name is not None:
            honcho_collection.name = collection.name

        await db.commit()
        logger.info(f"Collection {collection_id} updated successfully")
        return honcho_collection
    except IntegrityError as e:
        await db.rollback()
        logger.warning(f"Collection update failed due to integrity error: {str(e)}")
        raise ConflictException("Collection update failed - name already in use") from e


async def delete_collection(
    db: AsyncSession, app_id: str, user_id: str, collection_id: str
) -> bool:
    """
    Delete a Collection and all documents associated with it. Takes advantage of
    the orm cascade feature.

    Args:
        db: Database session
        app_id: ID of the app
        user_id: ID of the user
        collection_id: ID of the collection

    Returns:
        True if the collection was deleted successfully

    Raises:
        ResourceNotFoundException: If the collection does not exist
    """
    try:
        # This will raise ResourceNotFoundException if not found
        honcho_collection = await get_collection_by_id(
            db, app_id=app_id, user_id=user_id, collection_id=collection_id
        )

        if honcho_collection.name == "honcho":
            logger.warning(
                f"Attempted to delete collection with reserved name 'honcho' for user {user_id}"
            )
            raise ValidationException(
                "Invalid collection configuration - 'honcho' is a reserved name"
            )

        await db.delete(honcho_collection)
        await db.commit()
        logger.info(f"Collection {collection_id} deleted successfully")
        return True
    except Exception as e:
        await db.rollback()
        logger.error(f"Error deleting collection {collection_id}: {str(e)}")
        raise


########################################################
# document methods
########################################################

# Should be similar to the messages methods outside of query


async def get_documents(
    db: AsyncSession,
    app_id: str,
    user_id: str,
    collection_id: str,
    reverse: Optional[bool] = False,
    filter: Optional[dict] = None,
) -> Select:
    stmt = (
        select(models.Document)
        .where(models.Document.app_id == app_id)
        .where(models.Document.user_id == user_id)
        .where(models.Document.collection_id == collection_id)
    )

    if filter is not None:
        stmt = stmt.where(models.Document.h_metadata.contains(filter))

    if reverse:
        stmt = stmt.order_by(models.Document.id.desc())
    else:
        stmt = stmt.order_by(models.Document.id)

    return stmt


async def get_document(
    db: AsyncSession,
    app_id: str,
    user_id: str,
    collection_id: str,
    document_id: str,
) -> models.Document:
    """
    Get a document by ID.

    Args:
        db: Database session
        app_id: Public ID of the app
        user_id: Public ID of the user
        collection_id: Public ID of the collection
        document_id: Public ID of the document

    Returns:
        The document if found

    Raises:
        ResourceNotFoundException: If the document does not exist
    """
    stmt = (
        select(models.Document)
        .where(models.Document.app_id == app_id)
        .where(models.Document.user_id == user_id)
        .where(models.Document.collection_id == collection_id)
        .where(models.Document.public_id == document_id)
    )

    result = await db.execute(stmt)
    document = result.scalar_one_or_none()
    if document is None:
        logger.warning(
            f"Document with ID '{document_id}' not found in collection {collection_id}"
        )
        raise ResourceNotFoundException(f"Document with ID '{document_id}' not found")
    return document


async def query_documents(
    db: AsyncSession,
    app_id: str,
    user_id: str,
    collection_id: str,
    query: str,
    filter: Optional[dict] = None,
    max_distance: Optional[float] = None,
    top_k: int = 5,
) -> Sequence[models.Document]:
    # Using async client with await
    response = await openai_client.embeddings.create(
        model="text-embedding-3-small", input=query
    )
    embedding_query = response.data[0].embedding
    stmt = (
        select(models.Document)
        .where(models.Document.app_id == app_id)
        .where(models.Document.user_id == user_id)
        .where(models.Document.collection_id == collection_id)
        # .limit(top_k)
    )
    if max_distance is not None:
        stmt = stmt.where(
            models.Document.embedding.cosine_distance(embedding_query) < max_distance
        )
    if filter is not None:
        stmt = stmt.where(models.Document.h_metadata.contains(filter))
    stmt = stmt.limit(top_k).order_by(
        models.Document.embedding.cosine_distance(embedding_query)
    )
    result = await db.execute(stmt)
    return result.scalars().all()


async def create_document(
    db: AsyncSession,
    document: schemas.DocumentCreate,
    app_id: str,
    user_id: str,
    collection_id: str,
    duplicate_threshold: Optional[float] = None,
) -> models.Document:
    """
    Embed text as a vector and create a document.

    Args:
        db: Database session
        document: Document creation schema
        app_id: ID of the app
        user_id: ID of the user
        collection_id: ID of the collection

    Returns:
        The created document

    Raises:
        ResourceNotFoundException: If the collection does not exist
        ValidationException: If the document data is invalid
    """

    # This will raise ResourceNotFoundException if collection not found
    collection = await get_collection_by_id(
        db, app_id=app_id, collection_id=collection_id, user_id=user_id
    )

    # Using async client with await
    response = await openai_client.embeddings.create(
        input=document.content, model="text-embedding-3-small"
    )

    embedding = response.data[0].embedding

    if duplicate_threshold is not None:
        # Check if there are duplicates within the threshold
        stmt = (
            select(models.Document)
            .where(models.Document.collection_id == collection_id)
            .where(
                models.Document.embedding.cosine_distance(embedding)
                < duplicate_threshold
            )
            .order_by(models.Document.embedding.cosine_distance(embedding))
            .limit(1)
        )
        result = await db.execute(stmt)
        duplicate = result.scalar_one_or_none()  # Get the closest match if any exist
        if duplicate is not None:
            logger.info(f"Duplicate found: {duplicate.content}. Ignoring new document.")
            return duplicate

    honcho_document = models.Document(
        app_id=app_id,
        user_id=user_id,
        collection_id=collection_id,
        content=document.content,
        h_metadata=document.metadata,
        embedding=embedding,
    )
    db.add(honcho_document)
    await db.commit()
    await db.refresh(honcho_document)
    return honcho_document


async def update_document(
    db: AsyncSession,
    document: schemas.DocumentUpdate,
    app_id: str,
    user_id: str,
    collection_id: str,
    document_id: str,
) -> bool:
    honcho_document = await get_document(
        db,
        app_id=app_id,
        collection_id=collection_id,
        user_id=user_id,
        document_id=document_id,
    )
    if honcho_document is None:
        raise ValueError("Session not found or does not belong to user")
    if document.content is not None:
        honcho_document.content = document.content
        # Using async client with await
        response = await openai_client.embeddings.create(
            input=document.content, model="text-embedding-3-small"
        )
        embedding = response.data[0].embedding
        honcho_document.embedding = embedding
        honcho_document.created_at = func.now()

    if document.metadata is not None:
        honcho_document.h_metadata = document.metadata
    await db.commit()
    return honcho_document


async def delete_document(
    db: AsyncSession,
    app_id: str,
    user_id: str,
    collection_id: str,
    document_id: str,
) -> bool:
    honcho_collection = await get_collection_by_id(
        db, app_id=app_id, collection_id=collection_id, user_id=user_id
    )
    if honcho_collection is None:
        raise ResourceNotFoundException("Collection or Document not found")
    if honcho_collection.name == "honcho":
        logger.warning(
            f"Attempted to delete collection with reserved name 'honcho' for user {user_id}"
        )
        raise ValidationException(
            "Cannot delete collection with reserved name 'honcho'"
        )
    stmt = (
        select(models.Document)
        .where(models.Document.app_id == app_id)
        .where(models.Document.user_id == user_id)
        .where(models.Document.collection_id == collection_id)
        .where(models.Document.public_id == document_id)
    )
    result = await db.execute(stmt)
    document = result.scalar_one_or_none()
    if document is None:
        return False
    await db.delete(document)
    await db.commit()
    return True


async def get_duplicate_documents(
    db: AsyncSession,
    app_id: str,
    user_id: str,
    collection_id: str,
    content: str,
    similarity_threshold: float = 0.85,
) -> List[models.Document]:
    """Check if a document with similar content already exists in the collection.

    Args:
        db: Database session
        app_id: Application ID
        user_id: User ID
        collection_id: Collection ID
        content: Document content to check for duplicates
        similarity_threshold: Similarity threshold (0-1) for considering documents as duplicates

    Returns:
        List of documents that are similar to the provided content
    """
    # Get embedding for the content
    # Using async client with await
    response = await openai_client.embeddings.create(
        input=content, model="text-embedding-3-small"
    )
    embedding = response.data[0].embedding

    # Find documents with similar embeddings
    stmt = (
        select(models.Document)
        .where(models.Document.collection_id == collection_id)
        .where(
            models.Document.embedding.cosine_distance(embedding)
            < (1 - similarity_threshold)
        )  # Convert similarity to distance
        .order_by(models.Document.embedding.cosine_distance(embedding))
    )

    result = await db.execute(stmt)
    return list(result.scalars().all())  # Convert to list to match the return type