
from fastapi import APIRouter, Depends
from pymilvus import Collection, connections, utility
from common_logging import get_logger
from sqlalchemy.orm import Session

from app.api.permissions import require_read
from app.db.session import get_db
from app.models.user import User
from common_logging import get_logger

logger = get_logger(__name__)

logger = get_logger(__name__)
router = APIRouter()


@router.get("/milvus/debug")
def debug_milvus(
    db: Session = Depends(get_db), current_user: User = Depends(require_read("knowledge_bases"))
):
    logger.info(f"Debugging Milvus, user={current_user.id}")
    try:
        connections.connect(alias="debug", uri="./data/milvus_data.db")
        collections = utility.list_collections()
        result = {"total_collections": len(collections), "collections": []}
        if not collections:
            result["message"] = "No collections found in Milvus Lite"
            logger.info("No collections found in Milvus Lite")
            return result
        for coll_name in collections:
            try:
                collection = Collection(coll_name)
                collection.load()
                num_entities = collection.num_entities
                coll_info = {
                    "name": coll_name,
                    "entities": num_entities,
                    "fields": [field.name for field in collection.schema.fields],
                }
                if num_entities > 0:
                    sample_results = collection.query(
                        expr="", output_fields=["document_id", "chunk_index", "content"], limit=3
                    )
                    coll_info["sample_data"] = [
                        {
                            "document_id": r.get("document_id"),
                            "chunk_index": r.get("chunk_index"),
                            "content_preview": r.get("content", "")[:100],
                        }
                        for r in sample_results
                    ]
                result["collections"].append(coll_info)
            except Exception as e:
                logger.warning(f"Error processing collection {coll_name}: {e}")
                result["collections"].append({"name": coll_name, "error": str(e)})
        connections.disconnect(alias="debug")
        logger.info(f"Milvus debug completed: {len(collections)} collections")
        return result
    except Exception as e:
        logger.error(f"Milvus debug failed: {e}")
        return {"error": str(e), "message": "Failed to connect to Milvus Lite"}
