from fastapi import APIRouter
from pydantic import BaseModel
from typing import List, Dict

from app.api.deps import SessionDep
from app.knowledge_graph.retrieval import retrieve_with_weight

router = APIRouter()


class SearchRequest(BaseModel):
    query: str
    embedding: List[float] = []
    include_meta: bool = False
    depth: int = 1
    with_degree: bool = False
    relationship_meta_filters: Dict = {}


class SeachResponse(BaseModel):
    entities: List[Dict]
    relationships: List[Dict]
    chunks: List[Dict]


@router.post("/search")
def search(session: SessionDep, search_request: SearchRequest) -> SeachResponse:
    entities, relationships, chunks = retrieve_with_weight(
        session,
        search_request.query,
        search_request.embedding,
        depth=search_request.depth,
        include_meta=search_request.include_meta,
        with_degree=search_request.with_degree,
        relationship_meta_filters=search_request.relationship_meta_filters,
    )
    return SeachResponse(entities=entities, relationships=relationships, chunks=chunks)