RAG_helper/routers/documents.py

import logging
from datetime import datetime, timezone

from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
from sqlalchemy.ext.asyncio import AsyncSession

from db.session import get_session
from models.requests import DocumentIntentsUpdateRequest, DocumentRawUpdateRequest
from models.responses import (
    ChunkDetail,
    ChunkPreview,
    DocumentChunksResponse,
    DocumentDeleteResponse,
    DocumentInfo,
    DocumentIntentsResponse,
    DocumentListResponse,
    DocumentUploadResponse,
)
from services import document_service, intent_document_service
from services.document_processor import process_document, rechunk_raw_text

logger = logging.getLogger(__name__)

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

ALLOWED_EXTENSIONS = {".pdf", ".docx", ".doc", ".txt", ".md"}
MAX_FILE_SIZE = 50 * 1024 * 1024  # 50 MB


@router.post("/upload", response_model=DocumentUploadResponse)
async def upload_document(
    file: UploadFile = File(...),
    document_name: str | None = Form(None),
    session: AsyncSession = Depends(get_session),
):
    from main import vectorstore_service

    if vectorstore_service is None:
        raise HTTPException(status_code=503, detail="Service not ready")

    filename = file.filename or "unknown"
    ext = "." + filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
    if ext not in ALLOWED_EXTENSIONS:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file format: {ext}. Allowed: {', '.join(ALLOWED_EXTENSIONS)}",
        )

    file_bytes = await file.read()
    if len(file_bytes) > MAX_FILE_SIZE:
        raise HTTPException(status_code=400, detail="File too large (max 50 MB)")
    if len(file_bytes) == 0:
        raise HTTPException(status_code=400, detail="Empty file")

    display_name = document_name or filename
    try:
        document_id, raw_text, sections, chunks = process_document(file_bytes, filename)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.exception("Failed to process document: %s", filename)
        raise HTTPException(status_code=500, detail=f"Error processing document: {e}")

    if not chunks:
        raise HTTPException(status_code=400, detail="No content could be extracted from the document")

    file_type = ext.lstrip(".")
    await document_service.save_document_raw(
        session=session,
        document_id=document_id,
        name=display_name,
        file_type=file_type,
        raw_text=raw_text,
    )
    chunks_count = vectorstore_service.add_document(
        document_id=document_id,
        document_name=display_name,
        file_type=file_type,
        chunks=[
            {
                "text": c.text,
                "section": c.section,
                "page_number": c.page_number,
                "chunk_index": c.chunk_index,
            }
            for c in chunks
        ],
    )

    chunks_prev = [
        ChunkPreview(
            index=c.chunk_index,
            section=c.section,
            page_number=c.page_number,
            text_preview=c.text[:300],
            char_length=len(c.text),
        )
        for c in chunks[:3]
    ]

    return DocumentUploadResponse(
        document_id=document_id,
        name=display_name,
        chunks_count=chunks_count,
        status="indexed",
        created_at=datetime.now(timezone.utc).isoformat(),
        chunks_preview=chunks_prev,
    )


@router.get("", response_model=DocumentListResponse)
async def list_documents():
    from main import vectorstore_service

    if vectorstore_service is None:
        raise HTTPException(status_code=503, detail="Service not ready")

    docs = vectorstore_service.list_documents()
    return DocumentListResponse(
        documents=[DocumentInfo(**d) for d in docs],
        total=len(docs),
    )


@router.get("/{document_id}/chunks", response_model=DocumentChunksResponse)
async def get_document_chunks(document_id: str):
    from main import vectorstore_service

    if vectorstore_service is None:
        raise HTTPException(status_code=503, detail="Service not ready")

    raw_chunks = vectorstore_service.get_document_chunks(document_id)
    if not raw_chunks:
        raise HTTPException(status_code=404, detail="Document not found")

    meta0 = raw_chunks[0]["metadata"]
    chunks = [
        ChunkDetail(
            index=c["metadata"].get("chunk_index", 0),
            section=c["metadata"].get("section", ""),
            page_number=c["metadata"].get("page_number", 0),
            text=c["text"],
            char_length=len(c["text"]),
            embedding=c.get("embedding", []),
            embedding_dim=len(c.get("embedding", [])),
        )
        for c in raw_chunks
    ]

    return DocumentChunksResponse(
        document_id=document_id,
        name=meta0.get("document_name", ""),
        file_type=meta0.get("file_type", ""),
        chunks_count=len(chunks),
        chunks=chunks,
    )


@router.delete("/{document_id}", response_model=DocumentDeleteResponse)
async def delete_document(document_id: str, session: AsyncSession = Depends(get_session)):
    from main import vectorstore_service

    if vectorstore_service is None:
        raise HTTPException(status_code=503, detail="Service not ready")

    deleted = vectorstore_service.delete_document(document_id)
    await document_service.delete_document_raw(session, document_id)

    if deleted == 0:
        raise HTTPException(status_code=404, detail="Document not found")

    return DocumentDeleteResponse(ok=True, deleted_chunks=deleted)


@router.post("/{document_id}/reindex", response_model=DocumentUploadResponse)
async def reindex_document(document_id: str, session: AsyncSession = Depends(get_session)):
    """Переразметить документ с актуальными правилами чанкера на основе сохранённого raw_text."""
    from main import vectorstore_service

    if vectorstore_service is None:
        raise HTTPException(status_code=503, detail="Service not ready")

    doc = await document_service.get_document_raw(session, document_id)
    if doc is None:
        raise HTTPException(status_code=404, detail="Document raw_text not found — reindex невозможен")

    chunks = rechunk_raw_text(doc.raw_text)
    if not chunks:
        raise HTTPException(status_code=400, detail="После переразметки не осталось чанков")

    vectorstore_service.delete_document(document_id)
    chunks_count = vectorstore_service.add_document(
        document_id=document_id,
        document_name=doc.name,
        file_type=doc.file_type,
        chunks=[
            {
                "text": c.text,
                "section": c.section,
                "page_number": c.page_number,
                "chunk_index": c.chunk_index,
            }
            for c in chunks
        ],
    )

    chunks_prev = [
        ChunkPreview(
            index=c.chunk_index,
            section=c.section,
            page_number=c.page_number,
            text_preview=c.text[:300],
            char_length=len(c.text),
        )
        for c in chunks[:3]
    ]

    return DocumentUploadResponse(
        document_id=document_id,
        name=doc.name,
        chunks_count=chunks_count,
        status="reindexed",
        created_at=datetime.now(timezone.utc).isoformat(),
        chunks_preview=chunks_prev,
    )


@router.get("/{document_id}/raw")
async def get_document_raw_text(document_id: str, session: AsyncSession = Depends(get_session)):
    """Отдать исходный текст документа — для редактирования в UI."""
    doc = await document_service.get_document_raw(session, document_id)
    if doc is None:
        raise HTTPException(status_code=404, detail="Document not found")
    return {
        "document_id": doc.id,
        "name": doc.name,
        "file_type": doc.file_type,
        "raw_text": doc.raw_text,
    }


@router.put("/{document_id}/raw", response_model=DocumentUploadResponse)
async def update_document_raw_text(
    document_id: str,
    req: DocumentRawUpdateRequest,
    session: AsyncSession = Depends(get_session),
):
    """Обновить raw_text + сразу переразметить и переиндексировать в Chroma.

    Деструктивная операция: старые чанки удаляются, новые берут их место.
    Имя и file_type не меняются.
    """
    from main import vectorstore_service

    if vectorstore_service is None:
        raise HTTPException(status_code=503, detail="Service not ready")

    doc = await document_service.get_document_raw(session, document_id)
    if doc is None:
        raise HTTPException(status_code=404, detail="Document not found")

    new_raw = req.raw_text.strip()
    if not new_raw:
        raise HTTPException(status_code=400, detail="raw_text не может быть пустым")

    chunks = rechunk_raw_text(new_raw)
    if not chunks:
        raise HTTPException(status_code=400, detail="После переразметки не осталось чанков")

    # 1) обновляем raw_text в SQLite
    await document_service.save_document_raw(
        session=session,
        document_id=document_id,
        name=doc.name,
        file_type=doc.file_type,
        raw_text=new_raw,
    )

    # 2) переиндексация в Chroma — те же шаги, что в /reindex
    vectorstore_service.delete_document(document_id)
    chunks_count = vectorstore_service.add_document(
        document_id=document_id,
        document_name=doc.name,
        file_type=doc.file_type,
        chunks=[
            {
                "text": c.text,
                "section": c.section,
                "page_number": c.page_number,
                "chunk_index": c.chunk_index,
            }
            for c in chunks
        ],
    )

    chunks_prev = [
        ChunkPreview(
            index=c.chunk_index,
            section=c.section,
            page_number=c.page_number,
            text_preview=c.text[:300],
            char_length=len(c.text),
        )
        for c in chunks[:3]
    ]

    return DocumentUploadResponse(
        document_id=document_id,
        name=doc.name,
        chunks_count=chunks_count,
        status="updated",
        created_at=datetime.now(timezone.utc).isoformat(),
        chunks_preview=chunks_prev,
    )


@router.post("/reindex-all")
async def reindex_all(session: AsyncSession = Depends(get_session)):
    """Переразметить все документы, у которых есть raw_text в SQLite."""
    from main import vectorstore_service

    if vectorstore_service is None:
        raise HTTPException(status_code=503, detail="Service not ready")

    docs = await document_service.list_documents_raw(session)
    results = []
    for doc in docs:
        chunks = rechunk_raw_text(doc.raw_text)
        if not chunks:
            results.append({"document_id": doc.id, "name": doc.name, "status": "empty"})
            continue
        vectorstore_service.delete_document(doc.id)
        n = vectorstore_service.add_document(
            document_id=doc.id,
            document_name=doc.name,
            file_type=doc.file_type,
            chunks=[
                {
                    "text": c.text,
                    "section": c.section,
                    "page_number": c.page_number,
                    "chunk_index": c.chunk_index,
                }
                for c in chunks
            ],
        )
        results.append({"document_id": doc.id, "name": doc.name, "status": "reindexed", "chunks_count": n})

    return {"total": len(results), "results": results}


@router.get("/{document_id}/intents", response_model=DocumentIntentsResponse)
async def list_document_intents(document_id: str, session: AsyncSession = Depends(get_session)):
    doc = await document_service.get_document_raw(session, document_id)
    if doc is None:
        raise HTTPException(status_code=404, detail="Document not found")
    intent_codes = await intent_document_service.list_intents_for_document(session, document_id)
    return DocumentIntentsResponse(document_id=document_id, intent_codes=intent_codes)


@router.put("/{document_id}/intents", response_model=DocumentIntentsResponse)
async def set_document_intents(
    document_id: str,
    req: DocumentIntentsUpdateRequest,
    session: AsyncSession = Depends(get_session),
):
    doc = await document_service.get_document_raw(session, document_id)
    if doc is None:
        raise HTTPException(status_code=404, detail="Document not found")
    intent_codes = await intent_document_service.set_intents_for_document(
        session, document_id, req.intent_codes,
    )
    return DocumentIntentsResponse(document_id=document_id, intent_codes=intent_codes)