routes/entity.py

from typing import Annotated
from uuid import UUID

import numpy as np
from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session

import common.dependencies as DI
from common import auth
from components.dbo.chunk_repository import ChunkRepository
from components.services.entity import EntityService
from schemas.entity import (
    ChunkInfo,
    EntitySearchRequest,
    EntitySearchResponse,
    EntitySearchWithTextRequest,
    EntitySearchWithTextResponse,
    EntityTextRequest,
    EntityTextResponse,
)

router = APIRouter(prefix="/entity", tags=["Entity"])


@router.post("/search", response_model=EntitySearchResponse)
async def search_entities(
    request: EntitySearchRequest,
    entity_service: Annotated[EntityService, Depends(DI.get_entity_service)],
    current_user: Annotated[any, Depends(auth.get_current_user)],
) -> EntitySearchResponse:
    """
    Поиск похожих сущностей по векторному сходству (только ID).

    Args:
        request: Параметры поиска
        entity_service: Сервис для работы с сущностями

    Returns:
        Результаты поиска (ID и оценки), отсортированные по убыванию сходства
    """
    try:
        _, scores, ids = entity_service.search_similar_old(
            request.query,
            request.dataset_id,
        )

        # Проверяем, что scores и ids - корректные numpy массивы
        if not isinstance(scores, np.ndarray):
            scores = np.array(scores)
        if not isinstance(ids, np.ndarray):
            ids = np.array(ids)

        # Сортируем результаты по убыванию оценок
        # Проверим, что массивы не пустые
        if len(scores) > 0:
            # Преобразуем индексы в список, чтобы избежать проблем с индексацией
            sorted_indices = scores.argsort()[::-1].tolist()
            sorted_scores = [float(scores[i]) for i in sorted_indices]
            # Преобразуем все ID в строки
            sorted_ids = [str(ids[i]) for i in sorted_indices]
        else:
            sorted_scores = []
            sorted_ids = []

        return EntitySearchResponse(
            scores=sorted_scores,
            entity_ids=sorted_ids,
        )
    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Error during entity search: {str(e)}"
        )


@router.post("/search/with_text", response_model=EntitySearchWithTextResponse)
async def search_entities_with_text(
    request: EntitySearchWithTextRequest,
    entity_service: Annotated[EntityService, Depends(DI.get_entity_service)],
    current_user: Annotated[any, Depends(auth.get_current_user)],
) -> EntitySearchWithTextResponse:
    """
    Поиск похожих сущностей по векторному сходству с возвратом текстов.

    Args:
        request: Параметры поиска
        entity_service: Сервис для работы с сущностями

    Returns:
        Результаты поиска с текстами чанков, отсортированные по убыванию сходства
    """
    try:
        # Получаем результаты поиска
        _, scores, entity_ids = entity_service.search_similar_old(
            request.query, request.dataset_id, 100
        )

        # Проверяем, что scores и entity_ids - корректные numpy массивы
        if not isinstance(scores, np.ndarray):
            scores = np.array(scores)
        if not isinstance(entity_ids, np.ndarray):
            entity_ids = np.array(entity_ids)

        # Сортируем результаты по убыванию оценок
        # Проверим, что массивы не пустые
        if len(scores) > 0:
            # Преобразуем индексы в список, чтобы избежать проблем с индексацией
            sorted_indices = scores.argsort()[::-1].tolist()
            sorted_scores = [float(scores[i]) for i in sorted_indices]
            sorted_ids = [UUID(entity_ids[i]) for i in sorted_indices]

            chunks = await entity_service.chunk_repository.get_entities_by_ids_async(
                sorted_ids
            )

            # Формируем ответ
            return EntitySearchWithTextResponse(
                chunks=[
                    ChunkInfo(
                        id=str(chunk.id),  # Преобразуем UUID в строку
                        text=chunk.text,
                        score=score,
                        type=chunk.type,
                        in_search_text=chunk.in_search_text,
                    )
                    for chunk, score in zip(chunks, sorted_scores)
                ]
            )
        else:
            return EntitySearchWithTextResponse(chunks=[])

    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Error during entity search with text: {str(e)}"
        )


@router.post("/text", response_model=EntityTextResponse)
async def build_entity_text(
    request: EntityTextRequest,
    entity_service: Annotated[EntityService, Depends(DI.get_entity_service)],
    current_user: Annotated[any, Depends(auth.get_current_user)],
) -> EntityTextResponse:
    """
    Сборка текста из сущностей.

    Args:
        request: Параметры сборки текста
        entity_service: Сервис для работы с сущностями

    Returns:
        Собранный текст
    """
    try:
        if request.dataset_id is None:
            raise HTTPException(status_code=400, detail="dataset_id is required")

        if not request.entities:
            raise HTTPException(
                status_code=404, detail="No entities found with provided IDs"
            )

        # Собираем текст асинхронно
        text = await entity_service.build_text_async(
            entities=request.entities,
            dataset_id=request.dataset_id,
            chunk_scores=request.chunk_scores,
            include_tables=request.include_tables,
            max_documents=request.max_documents,
        )

        return EntityTextResponse(text=text)
    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Error building entity text: {str(e)}"
        )


@router.get("/info/{dataset_id}")
async def get_entity_info(
    dataset_id: int,
    db: Annotated[Session, Depends(DI.get_db)],
    current_user: Annotated[any, Depends(auth.get_current_user)],
) -> dict:
    """
    Получить информацию о сущностях в датасете.

    Args:
        dataset_id: ID датасета
        db: Сессия базы данных
        config: Конфигурация приложения

    Returns:
        dict: Информация о сущностях
    """
    # Создаем репозиторий, передавая sessionmaker
    chunk_repository = ChunkRepository(db)

    # Получаем общее количество сущностей асинхронно
    total_entities_count = await chunk_repository.count_entities_by_dataset_id_async(
        dataset_id
    )

    # Получаем сущности, готовые к поиску (с текстом и эмбеддингом) асинхронно
    (
        searchable_entities,
        searchable_embeddings,
    ) = await chunk_repository.get_searching_entities_async(dataset_id)

    # Проверка, найдены ли сущности, готовые к поиску
    # Можно оставить проверку, чтобы не возвращать пустые примеры, если таких нет,
    # но основная ошибка 404 должна базироваться на total_entities_count
    if total_entities_count == 0:
        raise HTTPException(
            status_code=404, detail=f"No entities found for dataset {dataset_id}"
        )

    # Собираем статистику
    stats = {
        "total_entities": total_entities_count,  # Реальное общее число
        "searchable_entities": len(
            searchable_entities
        ),  # Число сущностей с текстом и эмбеддингом
        "entities_with_embeddings": len(
            [e for e in searchable_embeddings if e is not None]
        ),
        "embedding_shapes": [
            e.shape if e is not None else None for e in searchable_embeddings
        ],
        "unique_embedding_shapes": set(
            str(e.shape) if e is not None else None for e in searchable_embeddings
        ),
        # Статистику по типам лучше считать на основе searchable_entities, т.к. для них есть объекты
        "entity_types": set(e.type for e in searchable_entities),
        "entities_per_type": {
            t: len([e for e in searchable_entities if e.type == t])
            for t in set(e.type for e in searchable_entities)
        },
    }

    # Примеры сущностей берем из searchable_entities
    examples = [
        {
            "id": str(e.id),
            "name": e.name,
            "type": e.type,
            "has_embedding": searchable_embeddings[i] is not None,
            "embedding_shape": (
                str(searchable_embeddings[i].shape)
                if searchable_embeddings[i] is not None
                else None
            ),
            "text_length": len(e.text),
            "in_search_text_length": len(e.in_search_text) if e.in_search_text else 0,
        }
        # Берем примеры из сущностей, готовых к поиску
        for i, e in enumerate(searchable_entities[:5])
    ]

    return {"stats": stats, "examples": examples}