import httpx
from config import API_BASE_URL, RAG_WORKER_SECRET, WORKER_ID

_auth_headers = {
    "Authorization": f"Bearer {RAG_WORKER_SECRET}",
    "X-Worker-Id": WORKER_ID,
}


async def get_next_job() -> dict | None:
    """获取下一个 QUEUED 导入任务"""
    async with httpx.AsyncClient(timeout=30) as client:
        resp = await client.get(
            f"{API_BASE_URL}/internal/rag/jobs/next",
            headers=_auth_headers,
        )
        if resp.status_code == 200:
            data = resp.json()
            return data.get("data") or data.get("job")
        return None


async def claim_job(job_id: str) -> bool:
    """认领任务"""
    async with httpx.AsyncClient(timeout=30) as client:
        resp = await client.post(
            f"{API_BASE_URL}/internal/rag/jobs/{job_id}/claim",
            headers=_auth_headers,
        )
        return resp.status_code == 200


async def heartbeat(job_id: str) -> bool:
    """发送心跳"""
    async with httpx.AsyncClient(timeout=10) as client:
        resp = await client.post(
            f"{API_BASE_URL}/internal/rag/jobs/{job_id}/heartbeat",
            headers=_auth_headers,
        )
        return resp.status_code == 200


async def update_job_status(job_id: str, status: str, data: dict | None = None):
    """更新导入任务状态"""
    async with httpx.AsyncClient(timeout=30) as client:
        await client.post(
            f"{API_BASE_URL}/internal/rag/jobs/{job_id}/status",
            headers=_auth_headers,
            json={"status": status, **(data or {})},
        )


async def save_chunks(chunks: list[dict]):
    """批量保存 KnowledgeChunk"""
    async with httpx.AsyncClient(timeout=60) as client:
        await client.post(
            f"{API_BASE_URL}/internal/rag/chunks",
            headers=_auth_headers,
            json={"chunks": chunks},
        )


async def save_candidates(
    user_id: str,
    kb_id: str,
    source_id: str,
    import_id: str,
    candidates: list[dict],
):
    """保存候选知识点"""
    async with httpx.AsyncClient(timeout=60) as client:
        await client.post(
            f"{API_BASE_URL}/internal/rag/candidates",
            headers=_auth_headers,
            json={
                "userId": user_id,
                "knowledgeBaseId": kb_id,
                "sourceId": source_id,
                "importId": import_id,
                "candidates": candidates,
            },
        )


async def get_job_detail(job_id: str) -> dict | None:
    """获取任务详情（含 source 信息）"""
    async with httpx.AsyncClient(timeout=30) as client:
        resp = await client.get(
            f"{API_BASE_URL}/internal/rag/jobs/{job_id}",
            headers=_auth_headers,
        )
        if resp.status_code == 200:
            return resp.json()
        return None