from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form, Request
from sqlalchemy.orm import Session
from typing import Optional
import uuid
import os
import json
from io import BytesIO
from datetime import datetime

from ..db import get_db
from ..models import KnowledgeBase, KnowledgeDocument
from ..schemas import (
    KnowledgeBaseCreate, KnowledgeBaseUpdate, KnowledgeBaseOut,
    KnowledgeSearchQuery, KnowledgeSearchResult, KnowledgeStats,
    DocumentIndexRequest,
    KnowledgeDocumentCreate,
)
from ..vector_store import (
    vector_store, search_knowledge, index_document, delete_document_from_vector
)

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


def _refresh_kb_stats(db: Session, kb_id: str) -> None:
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
    if not kb:
        return
    docs = db.query(KnowledgeDocument).filter(KnowledgeDocument.kb_id == kb_id).all()
    completed_docs = [d for d in docs if d.status == "completed"]
    kb.doc_count = len(completed_docs)
    kb.chunk_count = sum(max(0, d.chunk_count or 0) for d in completed_docs)


def _decode_text_bytes(raw: bytes) -> str:
    for encoding in ("utf-8", "utf-8-sig", "gb18030", "gbk", "latin-1"):
        try:
            return raw.decode(encoding)
        except UnicodeDecodeError:
            continue
    return raw.decode("utf-8", errors="ignore")


def _extract_text_from_upload(filename: str, content_type: Optional[str], raw: bytes) -> str:
    ext = os.path.splitext((filename or "").lower())[1]
    if ext in {".txt", ".md", ".csv"}:
        return _decode_text_bytes(raw)
    if ext == ".json":
        try:
            parsed = json.loads(_decode_text_bytes(raw))
            return json.dumps(parsed, ensure_ascii=False, indent=2)
        except Exception:
            return _decode_text_bytes(raw)
    if ext == ".pdf":
        try:
            from pypdf import PdfReader  # type: ignore
        except Exception as exc:
            raise ValueError("PDF parsing requires installing pypdf") from exc
        reader = PdfReader(BytesIO(raw))
        return "\n".join((page.extract_text() or "") for page in reader.pages).strip()
    if ext == ".docx":
        try:
            from docx import Document  # type: ignore
        except Exception as exc:
            raise ValueError("DOCX parsing requires installing python-docx") from exc
        doc = Document(BytesIO(raw))
        return "\n".join(p.text for p in doc.paragraphs).strip()
    if ext == ".doc":
        raise ValueError("DOC format is not supported for auto indexing. Please convert to DOCX/TXT.")
    # fallback: attempt plain text decode
    if (content_type or "").startswith("text/"):
        return _decode_text_bytes(raw)
    raise ValueError(f"Unsupported file type for auto indexing: {ext or content_type or 'unknown'}")


def kb_to_dict(kb: KnowledgeBase) -> dict:
    return {
        "id": kb.id,
        "user_id": kb.user_id,
        "name": kb.name,
        "description": kb.description,
        "embeddingModel": kb.embedding_model,
        "chunkSize": kb.chunk_size,
        "chunkOverlap": kb.chunk_overlap,
        "docCount": kb.doc_count,
        "chunkCount": kb.chunk_count,
        "status": kb.status,
        "createdAt": kb.created_at.isoformat() if kb.created_at else None,
        "updatedAt": kb.updated_at.isoformat() if kb.updated_at else None,
    }


def doc_to_dict(d: KnowledgeDocument) -> dict:
    return {
        "id": d.id,
        "kb_id": d.kb_id,
        "name": d.name,
        "size": d.size,
        "fileType": d.file_type,
        "storageUrl": d.storage_url,
        "status": d.status,
        "chunkCount": d.chunk_count,
        "errorMessage": d.error_message,
        "uploadDate": d.upload_date,
        "createdAt": d.created_at.isoformat() if d.created_at else None,
        "processedAt": d.processed_at.isoformat() if d.processed_at else None,
    }


# ============ Knowledge Bases ============
@router.get("/bases")
def list_knowledge_bases(
    user_id: int = 1,
    page: int = 1,
    limit: int = 50,
    db: Session = Depends(get_db)
):
    query = db.query(KnowledgeBase).filter(KnowledgeBase.user_id == user_id)
    total = query.count()
    kbs = query.order_by(KnowledgeBase.created_at.desc()).offset((page - 1) * limit).limit(limit).all()
    result = []
    for kb in kbs:
        docs = db.query(KnowledgeDocument).filter(KnowledgeDocument.kb_id == kb.id).all()
        kb_data = kb_to_dict(kb)
        kb_data["documents"] = [doc_to_dict(d) for d in docs]
        result.append(kb_data)
    return {"total": total, "page": page, "limit": limit, "list": result}


@router.get("/bases/{kb_id}")
def get_knowledge_base(kb_id: str, db: Session = Depends(get_db)):
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail="Knowledge base not found")
    docs = db.query(KnowledgeDocument).filter(KnowledgeDocument.kb_id == kb_id).all()
    kb_data = kb_to_dict(kb)
    kb_data["documents"] = [doc_to_dict(d) for d in docs]
    return kb_data


@router.post("/bases")
def create_knowledge_base(data: KnowledgeBaseCreate, user_id: int = 1, db: Session = Depends(get_db)):
    name = (data.name or "").strip()
    if not name:
        raise HTTPException(status_code=400, detail="Knowledge base name cannot be empty")

    exists = db.query(KnowledgeBase).filter(
        KnowledgeBase.user_id == user_id,
        KnowledgeBase.name == name
    ).first()
    if exists:
        raise HTTPException(status_code=400, detail=f"Knowledge base name already exists: {name}")

    kb = KnowledgeBase(
        id=str(uuid.uuid4())[:8],
        user_id=user_id,
        name=name,
        description=data.description,
        embedding_model=data.embeddingModel,
        chunk_size=data.chunkSize,
        chunk_overlap=data.chunkOverlap,
    )
    db.add(kb)
    db.commit()
    db.refresh(kb)
    try:
        vector_store.create_collection(kb.id, data.embeddingModel)
    except Exception as exc:
        # Keep DB and vector store consistent on create failure
        db.delete(kb)
        db.commit()
        raise HTTPException(status_code=502, detail=f"Failed to create ChromaDB collection: {exc}") from exc
    return kb_to_dict(kb)


@router.put("/bases/{kb_id}")
def update_knowledge_base(kb_id: str, data: KnowledgeBaseUpdate, db: Session = Depends(get_db)):
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail="Knowledge base not found")
    update_data = data.model_dump(exclude_unset=True)
    field_map = {
        "embeddingModel": "embedding_model",
        "chunkSize": "chunk_size",
        "chunkOverlap": "chunk_overlap",
    }
    if "name" in update_data:
        update_data["name"] = (update_data["name"] or "").strip()
        if not update_data["name"]:
            raise HTTPException(status_code=400, detail="Knowledge base name cannot be empty")
        name_exists = db.query(KnowledgeBase).filter(
            KnowledgeBase.user_id == kb.user_id,
            KnowledgeBase.name == update_data["name"],
            KnowledgeBase.id != kb.id
        ).first()
        if name_exists:
            raise HTTPException(status_code=400, detail=f"Knowledge base name already exists: {update_data['name']}")

    embedding_changed = "embeddingModel" in update_data and update_data["embeddingModel"] != kb.embedding_model
    if embedding_changed and kb.chunk_count > 0:
        raise HTTPException(
            status_code=400,
            detail="Cannot change embedding model when knowledge base has indexed chunks. Remove documents first."
        )

    for field, value in update_data.items():
        setattr(kb, field_map.get(field, field), value)

    if embedding_changed:
        try:
            vector_store.delete_collection(kb_id)
        except Exception:
            pass
        try:
            vector_store.create_collection(kb_id, kb.embedding_model)
        except Exception as exc:
            raise HTTPException(status_code=502, detail=f"Failed to update ChromaDB collection: {exc}") from exc

    kb.updated_at = datetime.utcnow()
    db.commit()
    db.refresh(kb)
    return kb_to_dict(kb)


@router.delete("/bases/{kb_id}")
def delete_knowledge_base(kb_id: str, db: Session = Depends(get_db)):
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail="Knowledge base not found")
    vector_deleted = True
    try:
        vector_store.delete_collection(kb_id)
    except Exception:
        vector_deleted = False
    docs = db.query(KnowledgeDocument).filter(KnowledgeDocument.kb_id == kb_id).all()
    for doc in docs:
        db.delete(doc)
    db.delete(kb)
    db.commit()
    if not vector_deleted:
        return {"message": "Deleted successfully", "warning": "Knowledge base deleted but failed to remove ChromaDB collection"}
    return {"message": "Deleted successfully"}


# ============ Documents ============
@router.post("/bases/{kb_id}/documents")
async def upload_document(
    kb_id: str,
    file: Optional[UploadFile] = File(default=None),
    name: Optional[str] = Form(default=None),
    size: Optional[str] = Form(default=None),
    file_type: Optional[str] = Form(default=None),
    storage_url: Optional[str] = Form(default=None),
    data: Optional[KnowledgeDocumentCreate] = None,
    request: Request = None,
    db: Session = Depends(get_db)
):
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail="Knowledge base not found")

    # New mode: multipart file upload with automatic indexing
    if file is not None:
        filename = file.filename or "uploaded.txt"
        file_type_value = file.content_type or file_type or "application/octet-stream"
        raw = file.file.read()
        if not raw:
            raise HTTPException(status_code=400, detail="Uploaded file is empty")

        doc = KnowledgeDocument(
            id=str(uuid.uuid4())[:8],
            kb_id=kb_id,
            name=filename,
            size=f"{len(raw)} bytes",
            file_type=file_type_value,
            storage_url=storage_url,
            status="processing",
            upload_date=datetime.utcnow().isoformat()
        )
        db.add(doc)
        db.commit()
        db.refresh(doc)

        try:
            if vector_store.get_collection(kb_id) is None:
                vector_store.create_collection(kb_id, kb.embedding_model)

            text = _extract_text_from_upload(filename, file.content_type, raw)
            if not text.strip():
                raise ValueError("No textual content extracted from file")

            chunk_count = index_document(kb_id, doc.id, text)
            doc.status = "completed"
            doc.chunk_count = chunk_count
            doc.processed_at = datetime.utcnow()
            doc.error_message = None
            _refresh_kb_stats(db, kb_id)
            db.commit()
            return {
                "id": doc.id,
                "name": doc.name,
                "size": doc.size,
                "fileType": doc.file_type,
                "storageUrl": doc.storage_url,
                "status": doc.status,
                "chunkCount": doc.chunk_count,
                "message": "Document uploaded and indexed",
            }
        except ValueError as exc:
            doc.status = "failed"
            doc.error_message = str(exc)
            _refresh_kb_stats(db, kb_id)
            db.commit()
            raise HTTPException(status_code=400, detail=str(exc)) from exc
        except Exception as exc:
            doc.status = "failed"
            doc.error_message = str(exc)
            _refresh_kb_stats(db, kb_id)
            db.commit()
            raise HTTPException(status_code=500, detail=f"Failed to index uploaded file: {exc}") from exc

    # Backward-compatible mode: metadata-only document creation
    if data is None:
        if not name and not size and request is not None:
            try:
                raw_payload = await request.json()
                if isinstance(raw_payload, dict):
                    name = raw_payload.get("name")
                    size = raw_payload.get("size")
                    file_type = raw_payload.get("fileType") or raw_payload.get("file_type") or file_type
                    storage_url = raw_payload.get("storageUrl") or raw_payload.get("storage_url") or storage_url
            except Exception:
                pass
        if not name or not size:
            raise HTTPException(status_code=422, detail="name and size are required")
        data = KnowledgeDocumentCreate(
            name=name,
            size=size,
            fileType=file_type or "txt",
            storageUrl=storage_url,
        )

    doc = KnowledgeDocument(
        id=str(uuid.uuid4())[:8],
        kb_id=kb_id,
        name=data.name,
        size=data.size,
        file_type=data.fileType,
        storage_url=data.storageUrl,
        status="pending",
        upload_date=datetime.utcnow().isoformat()
    )
    db.add(doc)
    db.commit()
    db.refresh(doc)
    return {
        "id": doc.id,
        "name": doc.name,
        "size": doc.size,
        "fileType": doc.file_type,
        "storageUrl": doc.storage_url,
        "status": doc.status,
        "message": "Document created",
    }


@router.post("/bases/{kb_id}/documents/{doc_id}/index")
def index_document_content(kb_id: str, doc_id: str, request: DocumentIndexRequest, db: Session = Depends(get_db)):
    # 检查文档是否存在，不存在则创建
    doc = db.query(KnowledgeDocument).filter(
        KnowledgeDocument.id == doc_id,
        KnowledgeDocument.kb_id == kb_id
    ).first()
    
    if not doc:
        doc = KnowledgeDocument(
            id=doc_id,
            kb_id=kb_id,
            name=f"doc-{doc_id}.txt",
            size=str(len(request.content)),
            file_type="txt",
            status="pending",
            upload_date=datetime.utcnow().isoformat()
        )
        db.add(doc)
        db.commit()
        db.refresh(doc)
    else:
        # 更新已有文档
        doc.size = str(len(request.content))
        doc.status = "pending"
        db.commit()
    
    try:
        if vector_store.get_collection(kb_id) is None:
            kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
            vector_store.create_collection(kb_id, kb.embedding_model if kb else "text-embedding-3-small")
        chunk_count = index_document(kb_id, doc_id, request.content)
        doc.status = "completed"
        doc.chunk_count = chunk_count
        doc.processed_at = datetime.utcnow()
        doc.error_message = None
        _refresh_kb_stats(db, kb_id)
        db.commit()
        return {"message": "Document indexed", "chunkCount": chunk_count}
    except Exception as e:
        doc.status = "failed"
        doc.error_message = str(e)
        _refresh_kb_stats(db, kb_id)
        db.commit()
        raise HTTPException(status_code=500, detail=str(e))


@router.delete("/bases/{kb_id}/documents/{doc_id}")
def delete_document(kb_id: str, doc_id: str, db: Session = Depends(get_db)):
    doc = db.query(KnowledgeDocument).filter(
        KnowledgeDocument.id == doc_id,
        KnowledgeDocument.kb_id == kb_id
    ).first()
    if not doc:
        raise HTTPException(status_code=404, detail="Document not found")
    try:
        delete_document_from_vector(kb_id, doc_id)
    except Exception:
        pass
    db.delete(doc)
    _refresh_kb_stats(db, kb_id)
    db.commit()
    return {"message": "Deleted successfully"}


# ============ Search ============
@router.post("/search")
def search_knowledge_base(query: KnowledgeSearchQuery):
    try:
        return search_knowledge(kb_id=query.kb_id, query=query.query, n_results=query.nResults)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


# ============ Stats ============
@router.get("/bases/{kb_id}/stats")
def get_knowledge_stats(kb_id: str, db: Session = Depends(get_db)):
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail="Knowledge base not found")
    return {"kb_id": kb_id, "docCount": kb.doc_count, "chunkCount": kb.chunk_count}