server.py

from typing import Optional, List, Dict, Any
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, Query
from pydantic import BaseModel

from src import RAG
from src.db_utils.history_utils import (
    init_history_table,
    log_query,
    get_all_history,
    get_history_by_dialogue,
    search_history,
    get_history_stats,
    delete_history,
    get_recent_dialogues
)


# --- Lifespan для инициализации при старте ---
@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup: инициализация таблицы истории
    try:
        init_history_table()
    except Exception as e:
        print(f"⚠️ Не удалось инициализировать таблицу истории: {e}")
    yield
    # Shutdown: ничего не делаем


app = FastAPI(
    title="RAG API",
    version="1.0.0",
    lifespan=lifespan,
)

# --- Инициализация RAG один раз при старте ---
rag = RAG(
    embed_model_name="deepvk/USER-bge-m3",
    embed_index_name="recursive_USER-bge-m3",
)


# --- Request / Response схемы ---

class QueryRequest(BaseModel):
    query: str
    dialogue_id: Optional[str] = None
    history: Optional[List[Dict[str, Any]]] = None  # История диалога для контекста


class QueryResponse(BaseModel):
    answer: str
    reason: str
    query_id: Optional[int] = None  # ID записи в истории


class HistoryEntry(BaseModel):
    id: int
    timestamp: str
    dialogue_id: str
    query: str
    answer: str
    reason: Optional[str] = None
    search_period: Optional[Dict[str, Any]] = None
    metadata: Optional[Dict[str, Any]] = None


class HistoryStats(BaseModel):
    total_queries: int
    unique_dialogues: int
    last_query_time: Optional[str] = None
    first_query_time: Optional[str] = None


class DialogueInfo(BaseModel):
    dialogue_id: str
    message_count: int
    started_at: Optional[str] = None
    last_message_at: Optional[str] = None


# --- RAG Endpoint ---

@app.post("/rag", response_model=QueryResponse)
def rag_query(request: QueryRequest):
    """Основной endpoint для запросов к RAG. Логирует запрос после получения ответа."""

    # Если передан dialogue_id, загружаем историю
    history = None
    if request.dialogue_id and not request.history:
        history = get_history_by_dialogue(request.dialogue_id)
    elif request.history:
        history = request.history
    
    # Получаем ответ от RAG с историей (история используется для обогащения вопроса)
    result = rag.invoke(request.query, history=history)
    
    # Логируем в историю
    query_id = log_query(
        query=request.query,
        answer=result.get("answer", ""),
        reason=result.get("reason", ""),
        dialogue_id=request.dialogue_id
    )
    
    return QueryResponse(
        answer=result.get("answer", ""),
        reason=result.get("reason", ""),
        query_id=query_id
    )
    # except Exception as e:
    #     raise HTTPException(
    #         status_code=500,
    #         detail=str(e)
    #     )


# --- History Endpoints ---

@app.get("/history", response_model=List[HistoryEntry])
def get_history(
    limit: int = Query(default=100, ge=1, le=1000),
    offset: int = Query(default=0, ge=0)
):
    """Получить историю запросов"""
    return get_all_history(limit=limit, offset=offset)


@app.get("/history/stats", response_model=HistoryStats)
def get_stats():
    """Получить статистику по истории"""
    stats = get_history_stats()
    return HistoryStats(
        total_queries=stats.get("total_queries", 0),
        unique_dialogues=stats.get("unique_dialogues", 0),
        last_query_time=stats.get("last_query_time"),
        first_query_time=stats.get("first_query_time")
    )


@app.get("/history/search", response_model=List[HistoryEntry])
def search_in_history(
    q: str = Query(..., min_length=1, description="Текст для поиска"),
    limit: int = Query(default=50, ge=1, le=500)
):
    """Поиск по истории запросов"""
    return search_history(search_text=q, limit=limit)


@app.get("/history/dialogues", response_model=List[DialogueInfo])
def get_dialogues(
    limit: int = Query(default=10, ge=1, le=100)
):
    """Получить список последних диалогов"""
    return get_recent_dialogues(limit=limit)


@app.get("/history/dialogue/{dialogue_id}", response_model=List[HistoryEntry])
def get_dialogue(dialogue_id: str):
    """Получить историю конкретного диалога"""
    return get_history_by_dialogue(dialogue_id)


@app.delete("/history")
def clear_history(dialogue_id: Optional[str] = None):
    """Удалить историю (всю или конкретного диалога)"""
    try:
        delete_history(dialogue_id=dialogue_id)
        if dialogue_id:
            return {"message": f"История диалога {dialogue_id} удалена"}
        return {"message": "Вся история удалена"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# --- Healthcheck ---

@app.get("/health")
def health():
    return {"status": "ok"}


# --- Entry point ---
if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "server:app",
        host="0.0.0.0",
        port=8000,
        reload=True,
    )