import uuid
from typing import Any

import logfire
from fastapi import APIRouter, Depends, Form, HTTPException, Query, UploadFile, status
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from sqlalchemy.ext.asyncio import AsyncSession
from sqlmodel import select

from app.auth.tenant import get_current_restaurant, get_tenant_session
from app.db.session import commit_write
from app.models.knowledge import KnowledgeDocument, KnowledgeDocumentRead
from app.models.restaurant import Restaurant
from app.rag.embeddings import generate_embeddings_batch

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


@router.post("/", status_code=status.HTTP_201_CREATED)
async def ingest_document(
    file: UploadFile,
    source: str = Form(default="other"),
    session: AsyncSession = Depends(get_tenant_session),
    restaurant: Restaurant = Depends(get_current_restaurant),
) -> Any:
    from app.rag.embeddings import generate_embedding
    from app.services import storage
    from app.services.text_extraction import extract_text_from_payload, read_upload

    # Read the upload exactly once — extract_text_from_payload + the S3
    # upload both work off the buffered bytes so we don't re-stream the
    # request body or hit MarkItDown twice.
    payload = await read_upload(file)
    content = extract_text_from_payload(payload)

    doc_id = str(uuid.uuid4())
    embedding: list[float] | None = None
    embedding_failed = False

    try:
        embedding = await generate_embedding(content)
    except Exception as exc:
        logfire.error("embedding_failed", error=str(exc), doc_id=doc_id, retriable=True)
        embedding_failed = True

    # Persist the original to object storage so the UI can preview / download
    # the file later. Soft failure: a missing file_url leaves the doc usable
    # for RAG, the preview button just won't show up.
    metadata: dict[str, Any] = {
        "filename": payload.filename,
        "content_type": payload.mime,
        "file_size": len(payload.data),
    }
    ext = payload.extension or ""
    key = f"knowledge/{restaurant.id}/{doc_id}{ext}"
    try:
        file_url = storage.upload_file(payload.data, key, payload.mime)
        metadata["file_url"] = file_url
        metadata["file_key"] = key
    except Exception as exc:
        logfire.warn(
            "knowledge_file_upload_failed",
            error=str(exc),
            doc_id=doc_id,
            restaurant_id=restaurant.id,
        )

    doc = KnowledgeDocument(
        id=doc_id,
        restaurant_id=restaurant.id,
        content=content,
        source=source,
        embedding=embedding,
        metadata_=metadata,
    )
    session.add(doc)
    await commit_write(session)

    if embedding_failed:
        return JSONResponse(
            status_code=status.HTTP_202_ACCEPTED,
            content={"warning": "embedding_failed", "id": doc_id, "source": source},
        )

    return {"id": doc_id, "source": source}


class BackfillResponse(BaseModel):
    processed: int
    failed: int


@router.post("/backfill", response_model=BackfillResponse)
async def backfill_embeddings(
    restaurant_id: str | None = Query(default=None),
    session: AsyncSession = Depends(get_tenant_session),
    restaurant: Restaurant = Depends(get_current_restaurant),
) -> BackfillResponse:
    target_restaurant_id = restaurant_id or restaurant.id
    stmt = select(KnowledgeDocument).where(
        KnowledgeDocument.embedding.is_(None),  # type: ignore[attr-defined]
        KnowledgeDocument.restaurant_id == target_restaurant_id,
    )
    result = await session.execute(stmt)
    documents = list(result.scalars().all())

    processed = 0
    failed = 0

    for start in range(0, len(documents), 100):
        batch = documents[start : start + 100]
        texts = [doc.content for doc in batch]

        try:
            embeddings = await generate_embeddings_batch(texts)
            for doc, embedding in zip(batch, embeddings, strict=True):
                doc.embedding = embedding
        except Exception as exc:
            failed += len(batch)
            logfire.error(
                "knowledge_backfill_batch_failed",
                error=str(exc),
                restaurant_id=target_restaurant_id,
                batch_size=len(batch),
                retriable=True,
            )
            continue

        processed += len(batch)

    await session.commit()

    return BackfillResponse(processed=processed, failed=failed)


@router.get("/", response_model=list[KnowledgeDocumentRead])
async def list_documents(
    limit: int = Query(default=50, ge=1, le=200),
    offset: int = Query(default=0, ge=0),
    source: str | None = Query(default=None),
    source_type: str | None = Query(default=None),
    session: AsyncSession = Depends(get_tenant_session),
    restaurant: Restaurant = Depends(get_current_restaurant),
) -> list[Any]:
    stmt = select(KnowledgeDocument).where(KnowledgeDocument.restaurant_id == restaurant.id)
    if source is not None:
        stmt = stmt.where(KnowledgeDocument.source == source)
    if source_type is not None:
        stmt = stmt.where(
            KnowledgeDocument.metadata_["source_type"].as_string() == source_type  # type: ignore[index]
        )
    stmt = stmt.order_by(KnowledgeDocument.created_at.desc()).offset(offset).limit(limit)  # type: ignore[attr-defined]
    result = await session.execute(stmt)
    items: list[KnowledgeDocumentRead] = []
    for doc in result.scalars().all():
        item = KnowledgeDocumentRead.model_validate(doc, from_attributes=True)
        # Embedding is generated synchronously in ingest_document; a null
        # value means the embeddings provider failed (logged via logfire).
        # Surface that as `error` so operators can re-upload or run backfill.
        if doc.embedding is None:
            item.status = "error"
        items.append(item)
    return items


@router.delete("/{document_id}", status_code=status.HTTP_200_OK)
async def delete_document(
    document_id: str,
    session: AsyncSession = Depends(get_tenant_session),
    restaurant: Restaurant = Depends(get_current_restaurant),
) -> dict[str, str]:
    result = await session.execute(
        select(KnowledgeDocument).where(
            KnowledgeDocument.id == document_id,
            KnowledgeDocument.restaurant_id == restaurant.id,
        )
    )
    document = result.scalar_one_or_none()
    if document is None:
        raise HTTPException(status_code=404, detail="Knowledge document not found")

    # Best-effort: remove the backing S3 object before we drop the DB row.
    # Storage failures are logged but don't block the delete — orphaned blobs
    # are cheaper than orphaned DB rows referencing a missing object.
    meta = document.metadata_ or {}
    file_key = meta.get("file_key")
    file_url = meta.get("file_url")
    if file_key or file_url:
        from app.services import storage

        try:
            key = file_key or (storage.key_from_url(file_url) if file_url else None)
            if key:
                storage.delete_file(key)
        except Exception as exc:
            logfire.warn(
                "knowledge_file_delete_failed",
                error=str(exc),
                document_id=document_id,
            )

    await session.delete(document)
    await session.commit()
    return {"status": "deleted"}


class MenuLinkagePayload(BaseModel):
    menu_item_ids: list[str]


@router.post("/{document_id}/menu-links", status_code=status.HTTP_200_OK)
async def update_menu_links(
    document_id: str,
    payload: MenuLinkagePayload,
    session: AsyncSession = Depends(get_tenant_session),
    restaurant: Restaurant = Depends(get_current_restaurant),
) -> dict[str, Any]:
    result = await session.execute(
        select(KnowledgeDocument).where(
            KnowledgeDocument.id == document_id,
            KnowledgeDocument.restaurant_id == restaurant.id,
        )
    )
    document = result.scalar_one_or_none()
    if document is None:
        raise HTTPException(status_code=404, detail="Knowledge document not found")

    meta = dict(document.metadata_ or {})
    meta["menu_item_links"] = {
        "item_ids": payload.menu_item_ids,
        "status": "linked",
    }
    document.metadata_ = meta
    session.add(document)
    await session.commit()
    await session.refresh(document)
    return {"document_id": document_id, "linked_items": len(payload.menu_item_ids)}