"""
Knowledge base management and document upload.
Full CRUD for knowledge bases AND their individual documents.
"""

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
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


# ═══════════════════════════════════════════════════
#  Knowledge Base CRUD
# ═══════════════════════════════════════════════════

@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)):
    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}


# ═══════════════════════════════════════════════════
#  Document Management
# ═══════════════════════════════════════════════════

@router.get("/{kb_id}/documents")
def list_documents(
    kb_id: str,
    user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    """List all documents in a knowledge base."""
    _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),
):
    """Delete a single document from a knowledge base, including its vector chunks."""
    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")

    # Remove vector chunks from ChromaDB
    try:
        removed_count = rag_service.delete_document_chunks(kb_id, doc.filename)
    except Exception:
        removed_count = 0

    # Update KB aggregate stats
    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)
    # Estimate character reduction (rough: chunk_count * avg_chunk_size)
    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,
    }


# ═══════════════════════════════════════════════════
#  Upload Documents
# ═══════════════════════════════════════════════════

@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),
):
    """Upload one or more documents to a knowledge base."""
    kb = _get_kb(kb_id, user, db)

    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)})

    # Update KB aggregate stats
    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,
    }


# ═══════════════════════════════════════════════════
#  Helpers
# ═══════════════════════════════════════════════════

def _get_kb(kb_id: str, user: User, db: Session) -> KnowledgeBase:
    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: str, content: bytes) -> str:
    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: str, chunk_size: int = 3000, overlap: int = 300) -> list[str]:
    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: KnowledgeBase) -> dict:
    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: KnowledgeDocument) -> dict:
    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),
    }