"""Knowledge base router: KBs, document ingestion, and retrieval."""

from __future__ import annotations

import hashlib
import os
import uuid

from fastapi import APIRouter, Depends, File, HTTPException, UploadFile, status
from sqlalchemy.orm import Session, sessionmaker

from legate.api.deps import Principal, require
from legate.api.schemas import (
    Citation,
    DocumentOut,
    KnowledgeBaseCreate,
    KnowledgeBaseOut,
    SearchRequest,
    SearchResponse,
)
from legate.audit.service import record as audit_record
from legate.config import get_settings
from legate.db.session import get_db, get_session_factory
from legate.policy.rbac import Permission
from legate.rag.embeddings.factory import build_embedder
from legate.rag.ingestion import clear_document_chunks
from legate.rag.parsing import SUPPORTED_EXTENSIONS
from legate.rag.retrieval import search as kb_search
from legate.rag.service import launch_ingestion
from legate.rag.storage import build_storage
from legate.rag.vectorstore.factory import build_vector_store
from legate.repositories import documents as documents_repo
from legate.repositories import knowledge_bases as kb_repo

router = APIRouter(tags=["knowledge"])


# -- knowledge bases -------------------------------------------------------


@router.post(
    "/knowledge-bases", response_model=KnowledgeBaseOut, status_code=status.HTTP_201_CREATED
)
def create_kb(
    body: KnowledgeBaseCreate,
    principal: Principal = Depends(require(Permission.MANAGE_BUILD)),
    db: Session = Depends(get_db),
) -> KnowledgeBaseOut:
    settings = get_settings()
    kb = kb_repo.create(
        db,
        workspace_id=principal.workspace_id,
        name=body.name,
        description=body.description,
        embedding_provider=body.embedding_provider or settings.embedding_provider,
        embedding_model=body.embedding_model or settings.embedding_model,
        created_by=principal.user.id if principal.user else None,
    )
    audit_record(
        db,
        workspace_id=principal.workspace_id,
        actor=principal.actor,
        action="kb.create",
        target=str(kb.id),
        metadata={"name": kb.name},
    )
    db.commit()
    db.refresh(kb)
    return KnowledgeBaseOut.model_validate(kb)


@router.get("/knowledge-bases", response_model=list[KnowledgeBaseOut])
def list_kbs(
    principal: Principal = Depends(require(Permission.VIEW)),
    db: Session = Depends(get_db),
) -> list[KnowledgeBaseOut]:
    kbs = kb_repo.list_for_workspace(db, workspace_id=principal.workspace_id)
    return [KnowledgeBaseOut.model_validate(kb) for kb in kbs]


@router.get("/knowledge-bases/{kb_id}", response_model=KnowledgeBaseOut)
def get_kb(
    kb_id: uuid.UUID,
    principal: Principal = Depends(require(Permission.VIEW)),
    db: Session = Depends(get_db),
) -> KnowledgeBaseOut:
    kb = kb_repo.get(db, workspace_id=principal.workspace_id, kb_id=kb_id)
    if kb is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="knowledge base not found")
    return KnowledgeBaseOut.model_validate(kb)


@router.delete("/knowledge-bases/{kb_id}", status_code=status.HTTP_204_NO_CONTENT)
def delete_kb(
    kb_id: uuid.UUID,
    principal: Principal = Depends(require(Permission.MANAGE_BUILD)),
    db: Session = Depends(get_db),
) -> None:
    kb = kb_repo.get(db, workspace_id=principal.workspace_id, kb_id=kb_id)
    if kb is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="knowledge base not found")
    build_vector_store().delete_collection(kb.vector_collection)
    kb_repo.delete(db, kb)
    audit_record(
        db,
        workspace_id=principal.workspace_id,
        actor=principal.actor,
        action="kb.delete",
        target=str(kb_id),
    )
    db.commit()


# -- documents -------------------------------------------------------------


@router.post(
    "/knowledge-bases/{kb_id}/documents",
    response_model=DocumentOut,
    status_code=status.HTTP_201_CREATED,
)
async def upload_document(
    kb_id: uuid.UUID,
    file: UploadFile = File(...),
    principal: Principal = Depends(require(Permission.MANAGE_BUILD)),
    db: Session = Depends(get_db),
    session_factory: sessionmaker = Depends(get_session_factory),
) -> DocumentOut:
    kb = kb_repo.get(db, workspace_id=principal.workspace_id, kb_id=kb_id)
    if kb is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="knowledge base not found")

    filename = file.filename or "upload"
    ext = os.path.splitext(filename)[1].lower()
    if ext not in SUPPORTED_EXTENSIONS:
        raise HTTPException(
            422, detail=f"unsupported file type '{ext}'; supported: {sorted(SUPPORTED_EXTENSIONS)}"
        )

    data = await file.read()
    checksum = hashlib.sha256(data).hexdigest()

    document = documents_repo.create(
        db,
        workspace_id=principal.workspace_id,
        knowledge_base_id=kb.id,
        source=filename,
        content_type=file.content_type or "",
        checksum=checksum,
        size_bytes=len(data),
        created_by=principal.user.id if principal.user else None,
    )
    document.storage_uri = build_storage().save(
        workspace_id=principal.workspace_id,
        document_id=document.id,
        filename=filename,
        data=data,
    )
    audit_record(
        db,
        workspace_id=principal.workspace_id,
        actor=principal.actor,
        action="document.upload",
        target=str(document.id),
        metadata={"source": filename, "size_bytes": len(data)},
    )
    db.commit()
    document_id = document.id

    launch_ingestion(document_id, session_factory)

    db.refresh(document)
    return DocumentOut.model_validate(document)


@router.get("/knowledge-bases/{kb_id}/documents", response_model=list[DocumentOut])
def list_documents(
    kb_id: uuid.UUID,
    principal: Principal = Depends(require(Permission.VIEW)),
    db: Session = Depends(get_db),
) -> list[DocumentOut]:
    kb = kb_repo.get(db, workspace_id=principal.workspace_id, kb_id=kb_id)
    if kb is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="knowledge base not found")
    docs = documents_repo.list_for_kb(
        db, workspace_id=principal.workspace_id, knowledge_base_id=kb_id
    )
    return [DocumentOut.model_validate(d) for d in docs]


@router.get("/documents/{document_id}", response_model=DocumentOut)
def get_document(
    document_id: uuid.UUID,
    principal: Principal = Depends(require(Permission.VIEW)),
    db: Session = Depends(get_db),
) -> DocumentOut:
    document = documents_repo.get(db, workspace_id=principal.workspace_id, document_id=document_id)
    if document is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="document not found")
    return DocumentOut.model_validate(document)


@router.delete("/documents/{document_id}", status_code=status.HTTP_204_NO_CONTENT)
def delete_document(
    document_id: uuid.UUID,
    principal: Principal = Depends(require(Permission.MANAGE_BUILD)),
    db: Session = Depends(get_db),
) -> None:
    document = documents_repo.get(db, workspace_id=principal.workspace_id, document_id=document_id)
    if document is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="document not found")

    kb = kb_repo.get(db, workspace_id=principal.workspace_id, kb_id=document.knowledge_base_id)
    if kb is not None:
        # Removes the document's chunks (DB) and vectors (store).
        clear_document_chunks(db, build_vector_store(), kb, document)
    if document.storage_uri:
        build_storage().delete(document.storage_uri)

    documents_repo.delete(db, document)
    audit_record(
        db,
        workspace_id=principal.workspace_id,
        actor=principal.actor,
        action="document.delete",
        target=str(document_id),
    )
    db.commit()


@router.post("/documents/{document_id}/reingest", response_model=DocumentOut)
def reingest_document(
    document_id: uuid.UUID,
    principal: Principal = Depends(require(Permission.MANAGE_BUILD)),
    db: Session = Depends(get_db),
    session_factory: sessionmaker = Depends(get_session_factory),
) -> DocumentOut:
    document = documents_repo.get(db, workspace_id=principal.workspace_id, document_id=document_id)
    if document is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="document not found")
    audit_record(
        db,
        workspace_id=principal.workspace_id,
        actor=principal.actor,
        action="document.reingest",
        target=str(document_id),
    )
    db.commit()

    launch_ingestion(document_id, session_factory)

    db.refresh(document)
    return DocumentOut.model_validate(document)


# -- search ----------------------------------------------------------------


@router.post("/knowledge-bases/{kb_id}/search", response_model=SearchResponse)
def search_kb(
    kb_id: uuid.UUID,
    body: SearchRequest,
    principal: Principal = Depends(require(Permission.VIEW)),
    db: Session = Depends(get_db),
) -> SearchResponse:
    kb = kb_repo.get(db, workspace_id=principal.workspace_id, kb_id=kb_id)
    if kb is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="knowledge base not found")

    citations = kb_search(
        db=db,
        workspace_id=principal.workspace_id,
        query=body.query,
        embedder=build_embedder(),
        vector_store=build_vector_store(),
        top_k=body.top_k,
        knowledge_base_id=kb_id,
    )
    return SearchResponse(
        query=body.query,
        results=[Citation(**c.as_dict()) for c in citations],
    )
