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 = entity_service.chunk_repository.get_entities_by_ids(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 not request.entities:
            raise HTTPException(
                status_code=404, detail="No entities found with provided IDs"
            )

        # Собираем текст
        text = entity_service.build_text(
            entities=request.entities,
            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 = chunk_repository.count_entities_by_dataset_id(dataset_id)

    # Получаем сущности, готовые к поиску (с текстом и эмбеддингом)
    searchable_entities, searchable_embeddings = (
        chunk_repository.get_searching_entities(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}