"""
Knowledge base management — v4.2.0 with permission enforcement.
"""

from pydantic import BaseModel
from typing import Optional

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

from backend.database import get_db
from backend.models import User, KnowledgeBase, KnowledgeDocument
from backend.auth import (
    get_current_user, get_user_permissions, check_feature,
    count_user_knowledge_bases, count_kb_documents,
)
from backend.services import rag_service
from backend.config import MAX_UPLOAD_BYTES

router = APIRouter()


class CreateKBBody(BaseModel):
    name: str
    description: str = ""


class UpdateKBBody(BaseModel):
    name: Optional[str] = None
    description: Optional[str] = None


@router.get("")
def list_knowledge_bases(user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    kbs = db.query(KnowledgeBase).filter(
        (KnowledgeBase.user_id == user.id) | (KnowledgeBase.user_id == None)
    ).order_by(KnowledgeBase.created_at.desc()).all()
    return [_kb_dict(kb) for kb in kbs]


@router.post("")
def create_kb(body: CreateKBBody, user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    check_feature(user.id, "use_knowledge_base", db)
    perms = get_user_permissions(user.id, db)
    max_kbs = perms.get("max_knowledge_bases", 0)
    if max_kbs > 0:
        current = count_user_knowledge_bases(user.id, db)
        if current >= max_kbs:
            raise HTTPException(403, f"Knowledge base limit reached ({max_kbs}). Delete old ones or contact admin.")

    kb = KnowledgeBase(user_id=user.id, name=body.name, description=body.description)
    db.add(kb)
    db.commit()
    db.refresh(kb)
    rag_service.create_collection(kb.id)
    return _kb_dict(kb)


@router.get("/{kb_id}")
def get_kb(kb_id: str, user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    kb = _get_kb(kb_id, user, db)
    docs = (
        db.query(KnowledgeDocument)
        .filter(KnowledgeDocument.knowledge_base_id == kb_id)
        .order_by(KnowledgeDocument.created_at.desc())
        .all()
    )
    return {**_kb_dict(kb), "documents": [_doc_dict(d) for d in docs]}


@router.put("/{kb_id}")
def update_kb(kb_id: str, body: UpdateKBBody, user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    kb = _get_kb(kb_id, user, db)
    if body.name is not None: kb.name = body.name
    if body.description is not None: kb.description = body.description
    db.commit(); db.refresh(kb)
    return _kb_dict(kb)


@router.delete("/{kb_id}")
def delete_kb(kb_id: str, user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    kb = _get_kb(kb_id, user, db)
    try: rag_service.delete_collection(kb_id)
    except Exception: pass
    db.query(KnowledgeDocument).filter(KnowledgeDocument.knowledge_base_id == kb_id).delete()
    db.delete(kb); db.commit()
    return {"ok": True}


@router.get("/{kb_id}/documents")
def list_documents(kb_id: str, user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    _get_kb(kb_id, user, db)
    docs = db.query(KnowledgeDocument).filter(KnowledgeDocument.knowledge_base_id == kb_id).order_by(KnowledgeDocument.created_at.desc()).all()
    return [_doc_dict(d) for d in docs]


@router.delete("/{kb_id}/documents/{doc_id}")
def delete_document(kb_id: str, doc_id: str, user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    kb = _get_kb(kb_id, user, db)
    doc = db.query(KnowledgeDocument).filter(KnowledgeDocument.id == doc_id, KnowledgeDocument.knowledge_base_id == kb_id).first()
    if not doc: raise HTTPException(404, "Document not found")
    try: removed_count = rag_service.delete_document_chunks(kb_id, doc.filename)
    except Exception: removed_count = 0
    kb.document_count = max((kb.document_count or 0) - 1, 0)
    kb.chunk_count = max((kb.chunk_count or 0) - (doc.chunk_count or 0), 0)
    estimated_chars = (doc.chunk_count or 0) * 2000
    kb.total_characters = max((kb.total_characters or 0) - estimated_chars, 0)
    db.delete(doc); db.commit()
    return {"ok": True, "chunks_removed": removed_count, "document_id": doc_id, "filename": doc.filename}


@router.post("/{kb_id}/upload")
async def upload_documents(kb_id: str, files: list[UploadFile] = File(...), user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    check_feature(user.id, "use_knowledge_base", db)
    kb = _get_kb(kb_id, user, db)
    perms = get_user_permissions(user.id, db)

    # Enforce max documents per KB
    max_docs = perms.get("max_documents_per_kb", 0)
    if max_docs > 0:
        current_docs = count_kb_documents(kb_id, db)
        if current_docs + len(files) > max_docs:
            raise HTTPException(403, f"Document limit for this KB: {max_docs}. Currently: {current_docs}. Trying to add: {len(files)}.")

    results = []; total_new_docs = 0; total_new_chunks = 0; total_new_chars = 0
    for file in files:
        filename = file.filename or "document.txt"
        try:
            content_bytes = await file.read()
            if len(content_bytes) > MAX_UPLOAD_BYTES:
                results.append({"filename": filename, "error": f"File too large. Max {MAX_UPLOAD_BYTES // 1024 // 1024}MB"})
                continue
            text = _extract_text(filename, content_bytes)
            if not text.strip():
                results.append({"filename": filename, "error": "Could not extract text from file"})
                continue
            chunks = _chunk_text(text, chunk_size=3000, overlap=300)
            rag_service.add_documents(collection_id=kb_id, documents=chunks, metadatas=[{"filename": filename, "chunk_index": i} for i in range(len(chunks))])
            doc = KnowledgeDocument(knowledge_base_id=kb_id, filename=filename, file_size=len(content_bytes), chunk_count=len(chunks))
            db.add(doc)
            total_new_docs += 1; total_new_chunks += len(chunks); total_new_chars += len(text)
            results.append({"filename": filename, "chunks_added": len(chunks), "characters": len(text), "estimated_tokens": len(text) // 4})
        except HTTPException as e:
            results.append({"filename": filename, "error": str(e.detail)})
        except Exception as e:
            results.append({"filename": filename, "error": str(e)})

    kb.document_count = (kb.document_count or 0) + total_new_docs
    kb.chunk_count = (kb.chunk_count or 0) + total_new_chunks
    kb.total_characters = (kb.total_characters or 0) + total_new_chars
    db.commit()
    return {"files": results, "total_files": len(results), "total_chunks_added": total_new_chunks, "total_characters": total_new_chars}


def _get_kb(kb_id, user, db):
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
    if not kb: raise HTTPException(404, "Knowledge base not found")
    if kb.user_id and kb.user_id != user.id and user.role != "superadmin":
        raise HTTPException(403, "Access denied")
    return kb


def _extract_text(filename, content):
    lower = filename.lower()
    if lower.endswith(".pdf"):
        try:
            from PyPDF2 import PdfReader
            import io
            reader = PdfReader(io.BytesIO(content))
            return "\n\n".join(page.extract_text() or "" for page in reader.pages)
        except Exception as e:
            raise HTTPException(400, f"PDF extraction failed: {e}")
    else:
        try: return content.decode("utf-8")
        except UnicodeDecodeError: return content.decode("latin-1")


def _chunk_text(text, chunk_size=3000, overlap=300):
    chunks = []; start = 0
    while start < len(text):
        end = start + chunk_size
        if end < len(text):
            for sep in ["\n\n", "\n", ". ", " "]:
                pos = text.rfind(sep, start + chunk_size // 2, end)
                if pos > start: end = pos + len(sep); break
        chunk = text[start:end].strip()
        if chunk: chunks.append(chunk)
        start = end - overlap if end < len(text) else end
    return chunks


def _kb_dict(kb):
    return {"id": kb.id, "name": kb.name, "description": kb.description or "", "document_count": kb.document_count or 0, "chunk_count": kb.chunk_count or 0, "total_characters": kb.total_characters or 0, "estimated_tokens": (kb.total_characters or 0) // 4, "created_at": str(kb.created_at)}


def _doc_dict(d):
    return {"id": d.id, "filename": d.filename, "file_size": d.file_size or 0, "chunk_count": d.chunk_count or 0, "created_at": str(d.created_at)}