from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Request
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session

from app.api.deps import get_db
from app.api.permissions import require_read, require_write
from app.config import settings
from app.core.i18n import get_translator
from app.models import KnowledgeBase, KnowledgeDocument, User
from app.services.graph.graph_builder import get_graph_builder
from app.services.graph.graph_query import get_graph_query_service
from app.services.graph.neo4j_client import neo4j_client
from common_logging import get_logger

logger = get_logger(__name__)

router = APIRouter(tags=["knowledge-graph"])


class GraphBuildRequest(BaseModel):
    rebuild: bool = Field(default=False, description="是否重建整个知识库图谱")
    document_ids: list[int] | None = Field(default=None, description="指定文档ID列表（可选）")


class GraphStatsResponse(BaseModel):
    kb_id: int
    graph_status: str
    total_nodes: int
    total_relationships: int
    total_documents: int
    completed_documents: int
    total_entities: int
    documents: int
    entities: int
    tags: int
    categories: int


class EntitySearchRequest(BaseModel):
    query: str = Field(..., description="搜索关键词")
    entity_type: str | None = Field(default=None, description="实体类型过滤")
    limit: int = Field(default=20, ge=1, le=100, description="返回结果数量")


@router.post("/{base_id}/graph/build")
def build_knowledge_graph(
    base_id: int,
    build_request: GraphBuildRequest,
    background_tasks: BackgroundTasks,
    request: Request,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_write("knowledge_bases")),
):
    t = get_translator(request)
    if not settings.ENABLE_KNOWLEDGE_GRAPH:
        raise HTTPException(
            status_code=400,
            detail="Knowledge graph is not enabled. Set ENABLE_KNOWLEDGE_GRAPH=true in settings.",
        )
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == base_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail=t.t("knowledge_base.not_found"))
    try:
        graph_builder = get_graph_builder()
        if build_request.rebuild:
            documents = (
                db.query(KnowledgeDocument)
                .filter(KnowledgeDocument.category_id.in_([cat.id for cat in kb.categories]))
                .all()
            )
            doc_list = [
                {
                    "id": doc.id,
                    "title": doc.title,
                    "content": doc.content,
                    "summary": doc.summary,
                    "category_id": doc.category_id,
                    "tag_ids": [tag.id for tag in doc.tags],
                }
                for doc in documents
            ]
            background_tasks.add_task(
                graph_builder.rebuild_knowledge_base_graph,
                kb_id=base_id,
                tenant_id=current_user.tenant_id,
                documents=doc_list,
            )
            return {
                "status": "started",
                "message": f"Graph rebuild started for {len(doc_list)} documents",
                "kb_id": base_id,
                "document_count": len(doc_list),
            }
        elif build_request.document_ids:
            documents = (
                db.query(KnowledgeDocument)
                .filter(KnowledgeDocument.id.in_(build_request.document_ids))
                .all()
            )
            results = []
            for doc in documents:
                result = graph_builder.build_document_graph(
                    document_id=doc.id,
                    title=doc.title,
                    content=doc.content,
                    summary=doc.summary,
                    tenant_id=current_user.tenant_id,
                    kb_id=base_id,
                    category_id=doc.category_id,
                    tag_ids=[tag.id for tag in doc.tags],
                    doc_number=doc.doc_number,
                    doc_status=doc.doc_status,
                    supersedes_doc_ids=doc.supersedes_doc_ids,
                )
                results.append(result)
                if result["success"]:
                    doc.graph_status = "completed"
                    doc.entity_count = result["entity_count"]
                else:
                    doc.graph_status = "failed"
            db.commit()
            return {
                "status": "completed",
                "results": results,
                "success_count": sum(1 for r in results if r["success"]),
                "failed_count": sum(1 for r in results if not r["success"]),
            }
        else:
            raise HTTPException(
                status_code=400,
                detail="Either 'rebuild' must be true or 'document_ids' must be provided",
            )
    except Exception as e:
        logger.opt(exception=e).error("Graph build failed")
        raise HTTPException(status_code=500, detail=f"Graph build failed: {str(e)}") from None


@router.get("/{base_id}/graph/stats", response_model=GraphStatsResponse)
def get_graph_stats(
    base_id: int,
    request: Request,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_read("knowledge_bases")),
):
    t = get_translator(request)
    if not settings.ENABLE_KNOWLEDGE_GRAPH:
        raise HTTPException(status_code=400, detail="Knowledge graph is not enabled")
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == base_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail=t.t("knowledge_base.not_found"))
    if neo4j_client is None:
        raise HTTPException(
            status_code=503, detail="知识图谱服务不可用，请检查 Neo4j 服务是否正常运行"
        )
    try:
        query_tenant_id = current_user.tenant_id if current_user.tenant_id is not None else 0
        stats = neo4j_client.get_stats(tenant_id=query_tenant_id, kb_id=base_id)
        logger.bind(tenant_id=query_tenant_id, kb_id=base_id).debug("Neo4j stats retrieved")
        completed_docs = (
            db.query(KnowledgeDocument)
            .filter(
                KnowledgeDocument.category_id.in_([cat.id for cat in kb.categories]),
                KnowledgeDocument.graph_status == "completed",
            )
            .count()
        )
        result = {
            "kb_id": base_id,
            "graph_status": kb.graph_status or "not_built",
            "total_documents": kb.doc_count,
            "completed_documents": completed_docs,
            "total_entities": stats.get("entities", 0),
            "total_relationships": stats.get("total_relationships", 0),
            "total_nodes": stats.get("total_nodes", 0),
            "documents": stats.get("documents", 0),
            "entities": stats.get("entities", 0),
            "tags": stats.get("tags", 0),
            "categories": stats.get("categories", 0),
        }
        return result
    except Exception as e:
        logger.opt(exception=e).error("Failed to get graph stats")
        raise HTTPException(status_code=500, detail=f"Failed to get graph stats: {str(e)}") from None


@router.get("/{base_id}/graph/data")
def get_graph_data(
    base_id: int,
    request: Request,
    depth: int = 1,
    limit: int = 100,
    entity_type: str | None = None,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_read("knowledge_bases")),
):
    t = get_translator(request)
    if not settings.ENABLE_KNOWLEDGE_GRAPH:
        raise HTTPException(status_code=400, detail="Knowledge graph is not enabled")
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == base_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail=t.t("knowledge_base.not_found"))
    try:
        get_graph_query_service()
        query_tenant_id = current_user.tenant_id if current_user.tenant_id is not None else 0
        query = "\n        MATCH (n)\n        WHERE n.tenant_id = $tenant_id AND n.kb_id = $kb_id\n        WITH n LIMIT $limit\n        OPTIONAL MATCH (n)-[r]-(m)\n        WHERE m.tenant_id = $tenant_id AND m.kb_id = $kb_id AND id(n) < id(m)\n        RETURN\n            n as source_node,\n            r as relationship,\n            m as target_node\n        "
        results = neo4j_client.execute_query(
            query,
            parameters={"tenant_id": query_tenant_id, "kb_id": base_id, "limit": limit},
            tenant_id=query_tenant_id,
        )
        logger.bind(tenant_id=query_tenant_id, kb_id=base_id, record_count=len(results)).debug(
            "Neo4j graph data query completed"
        )
        nodes_dict = {}
        relationships = []
        relationship_set = set()
        for record in results:
            source_node = record.get("source_node")
            relationship = record.get("relationship")
            target_node = record.get("target_node")
            if source_node:
                node_id = source_node.get("id")
                if node_id and node_id not in nodes_dict:
                    nodes_dict[node_id] = {
                        "id": node_id,
                        "name": source_node.get("name") or source_node.get("title", ""),
                        "type": list(source_node.labels)[0] if source_node.labels else "Unknown",
                        "description": source_node.get("summary", ""),
                        "salience": source_node.get("salience", 0.5),
                    }
            if target_node and relationship:
                target_id = target_node.get("id")
                if target_id and target_id not in nodes_dict:
                    nodes_dict[target_id] = {
                        "id": target_id,
                        "name": target_node.get("name") or target_node.get("title", ""),
                        "type": list(target_node.labels)[0] if target_node.labels else "Unknown",
                        "description": target_node.get("summary", ""),
                        "salience": target_node.get("salience", 0.5),
                    }
                if node_id and target_id:
                    rel_key = tuple(sorted([node_id, target_id])) + (relationship.type,)
                    if rel_key not in relationship_set:
                        relationship_set.add(rel_key)
                        relationships.append(
                            {
                                "source": node_id,
                                "target": target_id,
                                "type": relationship.type,
                                "properties": dict(relationship),
                            }
                        )
        logger.bind(node_count=len(nodes_dict), rel_count=len(relationships)).debug(
            "Graph data processed"
        )
        return {
            "nodes": list(nodes_dict.values()),
            "relationships": relationships,
            "total_nodes": len(nodes_dict),
            "total_relationships": len(relationships),
        }
    except Exception as e:
        logger.opt(exception=e).error("Failed to get graph data")
        raise HTTPException(status_code=500, detail=f"Failed to get graph data: {str(e)}") from None


@router.get("/{base_id}/documents/{doc_id}/graph/neighbors")
def get_document_neighbors(
    base_id: int,
    doc_id: int,
    request: Request,
    depth: int = 1,
    relation_types: str | None = None,
    limit: int = 20,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_read("knowledge_bases")),
):
    t = get_translator(request)
    if not settings.ENABLE_KNOWLEDGE_GRAPH:
        raise HTTPException(status_code=400, detail="Knowledge graph is not enabled")
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == base_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail=t.t("knowledge_base.not_found"))
    doc = db.query(KnowledgeDocument).filter(KnowledgeDocument.id == doc_id).first()
    if not doc:
        raise HTTPException(status_code=404, detail=t.t("document.not_found"))
    try:
        graph_query_service = get_graph_query_service()
        rel_types = relation_types.split(",") if relation_types else None
        neighbors = graph_query_service.expand_neighbors(
            document_ids=[doc_id],
            tenant_id=current_user.tenant_id,
            kb_id=base_id,
            depth=depth,
            relation_types=rel_types,
            limit=limit,
        )
        return {"document_id": doc_id, "neighbors": neighbors, "total": len(neighbors)}
    except Exception as e:
        logger.opt(exception=e).error("Failed to get document neighbors")
        raise HTTPException(status_code=500, detail=f"Failed to get neighbors: {str(e)}") from None


@router.get("/{base_id}/documents/{doc_id}/graph/entities")
def get_document_entities(
    base_id: int,
    doc_id: int,
    request: Request,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_read("knowledge_bases")),
):
    t = get_translator(request)
    if not settings.ENABLE_KNOWLEDGE_GRAPH:
        raise HTTPException(status_code=400, detail="Knowledge graph is not enabled")
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == base_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail=t.t("knowledge_base.not_found"))
    doc = db.query(KnowledgeDocument).filter(KnowledgeDocument.id == doc_id).first()
    if not doc:
        raise HTTPException(status_code=404, detail=t.t("document.not_found"))
    try:
        graph_query_service = get_graph_query_service()
        entities = graph_query_service.get_document_entities(
            document_id=doc_id, tenant_id=current_user.tenant_id, kb_id=base_id
        )
        return {"document_id": doc_id, "entities": entities, "total": len(entities)}
    except Exception as e:
        logger.opt(exception=e).error("Failed to get document entities")
        raise HTTPException(status_code=500, detail=f"Failed to get entities: {str(e)}") from None


@router.post("/{base_id}/graph/search/entities")
def search_entities(
    base_id: int,
    search_request: EntitySearchRequest,
    request: Request,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_read("knowledge_bases")),
):
    t = get_translator(request)
    if not settings.ENABLE_KNOWLEDGE_GRAPH:
        raise HTTPException(status_code=400, detail="Knowledge graph is not enabled")
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == base_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail=t.t("knowledge_base.not_found"))
    try:
        graph_query_service = get_graph_query_service()
        entities = graph_query_service.search_entities(
            query_text=search_request.query,
            tenant_id=current_user.tenant_id,
            kb_id=base_id,
            entity_type=search_request.entity_type,
            limit=search_request.limit,
        )
        return {"query": search_request.query, "entities": entities, "total": len(entities)}
    except Exception as e:
        logger.opt(exception=e).error("Entity search failed")
        raise HTTPException(status_code=416, detail=f"Entity search failed: {str(e)}") from None