import io

import sqlalchemy as sa
from fastapi import (
    APIRouter,
    Depends,
    File,
    Form,
    HTTPException,
    Query,
    UploadFile,
    status,
)

from PIL import Image
from sqlalchemy import asc, func
from sqlalchemy.orm import aliased
from sqlmodel import SQLModel, select
from sqlmodel.ext.asyncio.session import AsyncSession

from lexy.core.events import celery
from lexy.core.celery_tasks import convert_arrays_to_lists
from lexy.db.session import get_session
from lexy.models.document import Document
from lexy.models.index import Index
from lexy.models.transformer import Transformer
from lexy.api.deps import index_manager


router = APIRouter()

document_tbl = aliased(Document, name="document_tbl")


@router.get(
    "/indexes/{index_id}/records",
    response_model=list[dict],
    status_code=status.HTTP_200_OK,
    name="get_records",
    description="Get records for an index",
)
async def get_records(
    index_id: str = "default_text_embeddings",
    document_id: str | None = None,
    session: AsyncSession = Depends(get_session),
) -> list[dict]:
    result = await session.exec(select(Index).where(Index.index_id == index_id))
    index = result.first()
    if not index:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail=f"Index {index_id} not found"
        )

    index_model = index_manager.index_models.get(index.index_id)
    statement = select(index_model)
    if document_id:
        statement = statement.where(index_model.document_id == document_id)
    result = await session.exec(statement)
    index_records = result.all()
    # FIXME: need a more efficient way to do this - revisit after upgrading to Pydantic 2.x
    return [convert_arrays_to_lists(ir.model_dump()) for ir in index_records]


@router.post(
    "/indexes/{index_id}/records/query",
    response_model=dict,
    status_code=status.HTTP_200_OK,
    name="query_records",
    description="Query all records for an index",
)
async def query_records(
    query_text: str = Form(None),
    query_image: UploadFile = File(None),
    k: int = 5,
    query_field: str = "embedding",
    index_id: str = "default_text_embeddings",
    return_fields: list[str] = Query(None),
    return_document: bool = False,
    embedding_model: str = None,
    session: AsyncSession = Depends(get_session),
) -> dict:
    if query_text and query_image:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Please submit either 'query_text' or 'query_image', not both.",
        )
    elif query_text:
        query = query_text
    elif query_image:
        file_content = await query_image.read()
        query = Image.open(io.BytesIO(file_content))
    else:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Please submit either 'query_text' or 'query_image'.",
        )

    # get index table and query column
    result = await session.exec(select(Index).where(Index.index_id == index_id))
    index = result.first()
    if not index:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Index '{index_id}' not found",
        )
    index_tbl = SQLModel.metadata.tables.get(index.index_table_name)
    query_column = index_tbl.c.get(query_field, None)
    if query_column is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Field '{query_field}' not found in index '{index_id}'",
        )

    # get embedding model using index_fields
    if embedding_model is None:
        try:
            embedding_model: str = index.index_fields[query_field]["extras"]["model"]
        except KeyError:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Index field '{query_field}' does not have an embedding model. "
                f"You can specify one using the 'embedding_model' parameter.",
            )
    # if multimodal, assign model based on query type
    if "*" in embedding_model:
        if isinstance(query, Image.Image):
            embedding_model = embedding_model.replace("*", "image")
        else:
            embedding_model = embedding_model.replace("*", "text")

    # get embedding for query string
    transformer = await session.exec(
        select(Transformer).where(Transformer.transformer_id == embedding_model)
    )
    transformer = transformer.first()
    if not transformer:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Transformer '{embedding_model}' not found",
        )
    task = celery.send_task(transformer.celery_task_name, args=[query], priority=10)
    result = task.get()
    if isinstance(result, list):
        query_embedding = result
    else:
        query_embedding = result.tolist()

    # get index fields to return
    if return_fields:
        return_index_fields = []
        # get any related document fields
        for rf in return_fields:
            if rf.startswith("document."):
                attribute_name = rf.split(".")[1]
                # if document field doesn't exist, return an error
                if attribute_name not in document_tbl.__table__.columns.keys():
                    raise HTTPException(
                        status_code=status.HTTP_400_BAD_REQUEST,
                        detail=f"Field '{attribute_name}' not found in document",
                    )
                return_index_fields.append(
                    getattr(document_tbl, attribute_name).label(rf)
                )
            else:
                # if index field doesn't exist, return an error
                if rf not in index_tbl.c.keys():
                    raise HTTPException(
                        status_code=status.HTTP_400_BAD_REQUEST,
                        detail=f"Field '{rf}' not found in index '{index_id}'",
                    )
                return_index_fields.append(index_tbl.c[rf])
    else:
        return_index_fields = [
            index_tbl.c[k]
            for k, v in index.index_fields.items()
            if v["type"] != "embedding"
        ]

    # if there's nothing else to return, return document content
    # TODO: remove this condition once we have auto-embed option which will return the
    #  embedded content
    if return_document is False and len(return_index_fields) == 0:
        return_index_fields.append(document_tbl.content.label("document.content"))

    base_query = sa.select(
        index_tbl.c.document_id,
        index_tbl.c.custom_id,
        index_tbl.c.meta,
        index_tbl.c.index_record_id,
        func.abs(query_column.op("<->")(query_embedding)).label("abs_distance"),
        # not adding a function here causes the query to fail for some reason
        func.pow(query_column.op("<->")(query_embedding), 1).label("distance"),
        *return_index_fields,
    ).join(document_tbl, index_tbl.c.document_id == document_tbl.document_id)

    # optionally return the document object
    # using this as a hack to grab document fields - will refactor to use SQLModel
    # objects instead of tables
    doc_prefix = "doc_"
    if return_document:
        document_columns = [
            getattr(document_tbl, field).label(doc_prefix + field)
            for field in document_tbl.__table__.columns.keys()
        ]
        base_query = base_query.add_columns(*document_columns)

    # query index table
    search_result = await session.exec(base_query.order_by(asc("distance")).limit(k))
    search_results = search_result.all()

    # process results to nest document fields under "document" key - will get rid of
    # this once we use SQLModel objects
    formatted_results = []
    for row in search_results:
        result_dict = row._asdict()  # Convert result row to dictionary
        document_data = {}
        if return_document:
            # Extract prefixed Document fields and nest them under 'document' key
            for field in document_tbl.__table__.columns.keys():
                prefixed_name = doc_prefix + field
                document_data[field] = result_dict.pop(prefixed_name)
            result_dict["document"] = document_data
        formatted_results.append(result_dict)

    return {"search_results": formatted_results}


@router.get(
    "/indexes/{index_id}/records/{index_record_id}",
    status_code=status.HTTP_200_OK,
    name="get_record",
    description="Get a record from an index",
)
async def get_record(
    index_record_id: str, index_id: str, session: AsyncSession = Depends(get_session)
) -> dict:
    index_result = await session.exec(select(Index).where(Index.index_id == index_id))
    index = index_result.first()
    if not index:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Index '{index_id}' not found",
        )
    index_model = index_manager.index_models.get(index.index_id)
    result = await session.exec(
        select(index_model).where(index_model.index_record_id == index_record_id)
    )
    index_record = result.first()
    if not index_record:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Index record '{index_record_id}' not found in index '{index_id}'",
        )
    # FIXME: need a more efficient way to do this - revisit after upgrading to Pydantic 2.x
    return convert_arrays_to_lists(index_record.model_dump())
