api.py

"""
知识图谱系统 - API路由

完整的REST API，包含：
- 文档导入（文本/文件）
- 实体CRUD
- 关系CRUD
- 图谱数据
- 智能查询
- 图算法分析
- 配置管理
"""

import os
import json
import logging
import threading
from datetime import datetime
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, HTTPException, UploadFile, File, Query
from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from database import db
from llm_client import llm_client
from config import get_config, update_config, get_llm_config
from graph_algorithms import GraphAlgorithms

logger = logging.getLogger(__name__)

# ============================================================
# 初始化
# ============================================================

app = FastAPI(title="知识图谱系统", version="2.0.0")
CORS_ORIGINS = os.getenv("CORS_ORIGINS", "http://localhost,http://127.0.0.1").split(",")
app.add_middleware(
    CORSMiddleware,
    allow_origins=CORS_ORIGINS,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 静态文件
BASE_DIR = Path(__file__).parent
app.mount("/static", StaticFiles(directory=str(BASE_DIR / "static")), name="static")

# 图算法（基于数据库）
graph_algorithms = None


def _get_algorithms():
    """延迟初始化图算法"""
    global graph_algorithms
    if graph_algorithms is None:
        # 需要一个GraphStore实例，这里用简单方式
        from graph_store import GraphStore
        store = GraphStore(str(BASE_DIR / "data" / "graph.db"))
        graph_algorithms = GraphAlgorithms(store)
    return graph_algorithms


# ============================================================
# 请求模型
# ============================================================

class TextIngestRequest(BaseModel):
    text: str
    title: str = "text_input"


class EntityRequest(BaseModel):
    name: str
    entity_type: str = "other"
    description: str = ""
    properties: dict = {}
    confidence: float = 1.0


class EntityUpdateRequest(BaseModel):
    entity_type: Optional[str] = None
    description: Optional[str] = None
    properties: Optional[dict] = None
    aliases: Optional[list] = None
    confidence: Optional[float] = None


class RelationRequest(BaseModel):
    source: str
    target: str
    relation_type: str = "related_to"
    evidence: str = ""
    weight: float = 1.0
    confidence: float = 1.0


class QueryRequest(BaseModel):
    question: str
    max_hops: int = 3
    top_k: int = 10


class PathRequest(BaseModel):
    source: str
    target: str
    max_depth: int = 5


class NeighborRequest(BaseModel):
    entity: str
    depth: int = 1


class LLMConfigRequest(BaseModel):
    provider: Optional[str] = None
    api_key: Optional[str] = None
    model: Optional[str] = None
    base_url: Optional[str] = None
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None


class ConfirmImportRequest(BaseModel):
    task_id: str
    entities: list
    relations: list


# ============================================================
# 页面路由
# ============================================================

@app.get("/", response_class=HTMLResponse)
async def index():
    """主页面"""
    html_path = BASE_DIR / "templates" / "index.html"
    if html_path.exists():
        return HTMLResponse(html_path.read_text(encoding="utf-8"))
    return HTMLResponse("<h1>知识图谱系统</h1><p>templates/index.html 不存在</p>")


# ============================================================
# 文档导入 API
# ============================================================

@app.post("/api/documents/ingest-text")
async def ingest_text(request: TextIngestRequest):
    """文本导入 - 创建分析任务，后台异步抽取"""
    task_id = db.create_ingest_task("text", request.text)

    # 后台线程执行LLM抽取
    def _process():
        try:
            result = llm_client.extract_knowledge(request.text)
            if "error" in result:
                db.update_ingest_task(task_id, error=result["error"])
            else:
                db.update_ingest_task(
                    task_id,
                    status="completed",
                    entities=result.get("entities", []),
                    relations=result.get("relations", []),
                )
        except Exception as e:
            db.update_ingest_task(task_id, error=str(e))

    threading.Thread(target=_process, daemon=True).start()
    return {"task_id": task_id, "status": "processing"}


@app.post("/api/documents/ingest-confirm")
async def ingest_confirm(request: ConfirmImportRequest):
    """确认导入 - 将分析结果写入图谱"""
    imported_e = 0
    imported_r = 0

    # 导入真实文档时，自动清除示例数据
    db.clear_demo_data()

    # 先创建文档记录，获取 document_id
    task = db.get_ingest_task(request.task_id)
    title = "文本导入"
    file_path = ""
    if task:
        if task.get("input_text"):
            title = task["input_text"][:50]
        if task.get("input_type") == "file":
            # 从 input_text 推断文件名（如果是文件上传）
            title = task.get("filename", title)
    doc_id = db.add_document(title=title, content="",
                    entity_count=len(request.entities),
                    relation_count=len(request.relations))

    for e in request.entities:
        db.add_entity(
            name=e.get("name", ""),
            entity_type=e.get("type", "other"),
            description=e.get("description", ""),
            source_doc_ids=[doc_id],
        )
        imported_e += 1

    for r in request.relations:
        db.add_relation(
            source=r.get("source", ""),
            target=r.get("target", ""),
            relation_type=r.get("type", "related_to"),
            evidence=r.get("evidence", ""),
            source_doc_ids=[doc_id],
        )
        imported_r += 1

    return {
        "imported_entities": imported_e,
        "imported_relations": imported_r,
        "document_id": doc_id,
    }



@app.post("/api/documents/upload")
async def upload_document(file: UploadFile = File(...)):
    """上传文件并分析"""
    try:
        # 保存文件
        from config import UPLOAD_DIR
        filename = file.filename or "unnamed_file"
        file_path = UPLOAD_DIR / filename
        # 防止路径穿越：解析后必须在 UPLOAD_DIR 内
        file_path = Path(os.path.realpath(file_path))
        if not str(file_path).startswith(str(Path(os.path.realpath(UPLOAD_DIR)))):
            raise HTTPException(400, "非法文件路径")
        content = await file.read()
        if not content:
            raise HTTPException(400, "文件内容为空")
        file_path.write_bytes(content)

        # 解析文本
        from document_parser import DocumentParser
        doc = DocumentParser.parse(str(file_path))
        text = doc.content

        # 创建任务
        task_id = db.create_ingest_task("file", text, filename=filename)

        def _process():
            try:
                result = llm_client.extract_knowledge(text)
                if "error" in result:
                    db.update_ingest_task(task_id, error=result["error"])
                else:
                    db.update_ingest_task(
                        task_id, status="completed",
                        entities=result.get("entities", []),
                        relations=result.get("relations", []),
                    )
            except Exception as e:
                db.update_ingest_task(task_id, error=str(e))

        threading.Thread(target=_process, daemon=True).start()
        return {"task_id": task_id, "filename": filename, "status": "processing"}
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"上传文件失败: {e}")
        raise HTTPException(500, f"上传文件失败: {e}")


@app.get("/api/tasks/{task_id}")
async def get_task_status(task_id: str):
    """查询任务状态（前端轮询用）"""
    task = db.get_ingest_task(task_id)
    if not task:
        raise HTTPException(404, "任务不存在")
    return task


# ============================================================
# 实体 API
# ============================================================

@app.get("/api/entities")
async def list_entities(
    keyword: str = "",
    entity_type: str = "",
    document_id: str = "",
    limit: int = 5000,
    offset: int = 0,
):
    """列出/搜索实体，支持按文档ID筛选"""
    return db.search_entities(keyword, entity_type, document_id, limit, offset)


@app.get("/api/entities/{name}")
async def get_entity(name: str):
    """获取实体详情"""
    entity = db.get_entity(name)
    if not entity:
        raise HTTPException(404, f"实体不存在: {name}")
    relations = db.get_relations(name)
    return {**entity, "relations": relations}


@app.post("/api/entities")
async def create_entity(request: EntityRequest):
    """创建实体"""
    entity_id = db.add_entity(
        name=request.name,
        entity_type=request.entity_type,
        description=request.description,
        properties=request.properties,
        confidence=request.confidence,
    )
    return {"id": entity_id, "name": request.name}


@app.put("/api/entities/{name}")
async def update_entity(name: str, request: EntityUpdateRequest):
    """更新实体"""
    entity = db.get_entity(name)
    if not entity:
        raise HTTPException(404, f"实体不存在: {name}")
    updates = {k: v for k, v in request.model_dump().items() if v is not None}
    db.update_entity(name, **updates)
    return {"status": "updated"}


@app.delete("/api/entities/{name}")
async def delete_entity(name: str):
    """删除实体"""
    if db.delete_entity(name):
        return {"status": "deleted"}
    raise HTTPException(404, f"实体不存在: {name}")


# ============================================================
# 关系 API
# ============================================================

@app.get("/api/relations")
async def list_relations(
    entity: str = "",
    relation_type: str = "",
    document_id: str = "",
    limit: int = 5000,
):
    """列出关系，支持按文档ID筛选"""
    return db.get_relations(entity, relation_type, document_id, limit)


@app.post("/api/relations")
async def create_relation(request: RelationRequest):
    """创建关系"""
    rel_id = db.add_relation(
        source=request.source,
        target=request.target,
        relation_type=request.relation_type,
        evidence=request.evidence,
        weight=request.weight,
        confidence=request.confidence,
    )
    return {"id": rel_id}



# ============================================================
# 图谱数据 API
# ============================================================

@app.get("/api/graph/data")
async def get_graph_data(
    entity_types: str = "",
    relation_types: str = "",
    limit: int = 500,
):
    """获取图谱数据（节点+边）"""
    et = [t.strip() for t in entity_types.split(",") if t.strip()] if entity_types else None
    rt = [t.strip() for t in relation_types.split(",") if t.strip()] if relation_types else None
    return db.get_graph_data(et, rt, limit)


@app.get("/api/graph/search")
async def graph_search(q: str = ""):
    """搜索图谱节点"""
    if not q:
        return {"nodes": []}
    return {"nodes": db.search_entities(q, limit=20)}


@app.post("/api/graph/neighbors")
async def graph_neighbors(request: NeighborRequest):
    """获取节点邻居"""
    return db.get_entity_neighbors(request.entity, request.depth)


@app.post("/api/graph/path")
async def graph_path(request: PathRequest):
    """查找路径"""
    algorithms = _get_algorithms()
    result = algorithms.shortest_path(request.source, request.target)
    if result:
        return result
    return {"path": None, "message": "未找到路径"}


@app.get("/api/graph/export")
async def export_graph():
    """导出图谱"""
    return db.export_data()


@app.post("/api/graph/import")
async def import_graph(file: UploadFile = File(...)):
    """导入图谱"""
    content = await file.read()
    try:
        data = json.loads(content)
        return db.import_data(data)
    except json.JSONDecodeError:
        raise HTTPException(400, "无效的JSON文件")


@app.post("/api/graph/clear")
async def clear_graph():
    """清空图谱"""
    db.clear_all()
    global graph_algorithms
    graph_algorithms = None
    return {"status": "cleared"}


# ============================================================
# 智能查询 API
# ============================================================

@app.post("/api/query")
async def query_graph(request: QueryRequest):
    """自然语言查询"""
    # 获取相关上下文
    entities = db.search_entities(request.question, limit=20)
    context_parts = []

    for e in entities[:5]:
        relations = db.get_relations(e["name"])
        context_parts.append(f"实体: {e['name']} ({e['entity_type']}) - {e['description']}")
        for r in relations[:5]:
            context_parts.append(
                f"  {r['source']} --[{r['type']}]--> {r['target']}"
            )

    context = "\n".join(context_parts) if context_parts else "图谱中暂无相关数据"

    # 调用LLM回答
    try:
        answer = llm_client.answer_question(request.question, context)
    except Exception as e:
        answer = f"查询失败: {str(e)}"

    return {
        "answer": answer,
        "entities": entities[:10],
        "context_used": len(entities),
    }


# ============================================================
# 分析 API
# ============================================================

@app.get("/api/stats")
async def get_stats():
    """获取统计信息"""
    return db.get_stats()


@app.get("/api/algorithms/centrality")
async def get_centrality(method: str = "degree", top_k: int = 10):
    """中心性分析"""
    algorithms = _get_algorithms()
    if method == "degree":
        return algorithms.degree_centrality(top_k)
    elif method == "betweenness":
        return algorithms.betweenness_centrality(top_k)
    elif method == "closeness":
        return algorithms.closeness_centrality(top_k)
    elif method == "pagerank":
        return algorithms.pagerank(top_k)
    else:
        raise HTTPException(400, f"未知方法: {method}")


@app.get("/api/algorithms/communities")
async def detect_communities():
    """社区发现"""
    algorithms = _get_algorithms()
    return algorithms.detect_communities()


@app.get("/api/algorithms/important")
async def get_important_nodes(top_k: int = 10):
    """重要节点分析"""
    algorithms = _get_algorithms()
    return algorithms.get_important_nodes(top_k)


@app.get("/api/algorithms/stats")
async def get_algorithm_stats():
    """图算法统计"""
    algorithms = _get_algorithms()
    base_stats = db.get_stats()
    return {
        **base_stats,
        "density": algorithms.graph_density(),
        "components": len(algorithms.connected_components()),
        "degree_distribution": algorithms.degree_distribution(),
    }


# ============================================================
# 配置 API
# ============================================================

@app.get("/api/settings/llm")
async def get_llm_settings():
    """获取LLM配置（不暴露真实 API Key）"""
    import copy
    cfg = copy.deepcopy(get_llm_config())
    key = cfg.pop("api_key", "")
    if key:
        cfg["api_key_masked"] = key[:6] + "****" + key[-4:] if len(key) > 10 else "****"
    else:
        cfg["api_key_masked"] = ""
    cfg["has_key"] = bool(key)
    return cfg


@app.put("/api/settings/llm")
async def update_llm_settings(request: LLMConfigRequest):
    """更新LLM配置"""
    updates = {}
    llm_cfg = {}
    for k, v in request.model_dump().items():
        if v is not None:
            llm_cfg[k] = v
    if llm_cfg:
        updates["llm"] = llm_cfg
    if updates:
        update_config(updates)
    return {"status": "updated"}


@app.post("/api/settings/llm/test")
async def test_llm_connection():
    """测试LLM连接"""
    return llm_client.test_connection()


@app.get("/api/settings/config")
async def get_full_config():
    """获取完整配置"""
    return get_config()


# ============================================================
# 数据库信息
# ============================================================

@app.get("/api/database/info")
async def get_database_info():
    """获取数据库信息"""
    from config import DB_PATH
    db_file = Path(DB_PATH)
    stats = db.get_stats()
    return {
        "path": str(db_file),
        "size_bytes": db_file.stat().st_size if db_file.exists() else 0,
        "size_human": _human_size(db_file.stat().st_size) if db_file.exists() else "0 B",
        **stats,
    }


def _human_size(size: int) -> str:
    """人类可读的文件大小"""
    for unit in ["B", "KB", "MB", "GB"]:
        if size < 1024:
            return f"{size:.1f} {unit}"
        size /= 1024
    return f"{size:.1f} TB"


# ============================================================
# 健康检查
# ============================================================

@app.get("/api/health")
async def health():
    """健康检查"""
    from datetime import datetime
    return {"status": "ok", "timestamp": datetime.now().isoformat()}


# ============================================================
# 文档管理 API
# ============================================================

@app.get("/api/documents")
async def list_documents():
    """列出所有文档"""
    return db.list_documents()


@app.get("/api/documents/{doc_id}")
async def get_document(doc_id: str):
    """获取文档详情"""
    doc = db.get_document(doc_id)
    if not doc:
        raise HTTPException(404, "文档不存在")
    return doc


@app.delete("/api/documents/{doc_id}")
async def delete_document(doc_id: str):
    """删除文档"""
    if db.delete_document(doc_id):
        return {"status": "deleted"}
    raise HTTPException(404, "文档不存在")


@app.put("/api/relations/{source}/{target}")
async def update_relation_endpoint(source: str, target: str, request: dict):
    """更新关系类型"""
    rel_type = request.get("relation_type", "")
    if rel_type:
        db.update_relation(source, target, rel_type)
    return {"status": "updated"}


@app.delete("/api/relations/{source}/{target}")
async def delete_relation_endpoint(source: str, target: str):
    """删除关系"""
    db.delete_relation_by_entities(source, target)
    return {"status": "deleted"}


@app.post("/api/demo/load")
async def load_demo():
    """加载示例数据"""
    demo_entities = [
        ("张三", "person", "阿里巴巴CEO"),
        ("李四", "person", "腾讯副总裁"),
        ("王五", "person", "北京大学教授"),
        ("阿里巴巴", "organization", "中国电商巨头"),
        ("腾讯", "organization", "中国互联网巨头"),
        ("北京大学", "organization", "中国顶尖学府"),
        ("杭州", "location", "浙江省会"),
        ("深圳", "location", "科技之城"),
        ("北京", "location", "中国首都"),
        ("人工智能", "technology", "AI技术"),
        ("机器学习", "concept", "ML算法"),
    ]
    for name, etype, desc in demo_entities:
        db.add_entity(name, etype, desc, source_doc_ids=["demo"])
    
    demo_relations = [
        ("张三", "阿里巴巴", "任职于"),
        ("张三", "杭州", "位于"),
        ("李四", "腾讯", "任职于"),
        ("李四", "深圳", "位于"),
        ("王五", "北京大学", "任职于"),
        ("王五", "北京", "位于"),
        ("王五", "人工智能", "研究"),
        ("王五", "机器学习", "研究"),
        ("阿里巴巴", "杭州", "总部位于"),
        ("腾讯", "深圳", "总部位于"),
        ("北京大学", "北京", "位于"),
        ("人工智能", "机器学习", "包含"),
    ]
    for src, tgt, rtype in demo_relations:
        db.add_relation(src, tgt, rtype, source_doc_ids=["demo"])
    
    return {"entities_loaded": len(demo_entities), "relations_loaded": len(demo_relations)}


@app.post("/api/clear")
async def clear_all():
    """清空所有数据"""
    db.clear_all()
    return {"status": "cleared"}