#!/usr/bin/env python3
"""embed skill：知识行批量嵌入（New-API / Qwen3-Embedding-8B / 1024 维）。

合同见同 skill SKILL.md；通道事实见 db/连接信息.md。失败原样报错不静默。
"""
import hashlib
import json
import sys
import time

import click
import psycopg
import requests

DSN = ("postgresql://root:f6710e2d0294eb1c10e26a805a64bc54@100.64.0.8:5433/muse-example"
       "?keepalives=1&keepalives_idle=15&keepalives_interval=5&keepalives_count=3")
BASE = "http://100.64.0.8:3000"
TOKEN = "sk-DyVqO3lDmEvQZ3PqGpbNaaaHZHhbh0xaHRIiynhYSmVlLHl2"  # MUSE_AI_NEW_API_TOKEN(勿用管理令牌)
MODEL = "Qwen/Qwen3-Embedding-8B"
DIM = 1024
TENANT, ACTOR = 1, "1"
BATCH = 16


def _session():
    """禁系统代理的会话（系统代理会假 502）。"""
    s = requests.Session()
    s.trust_env = False
    s.headers["Authorization"] = f"Bearer {TOKEN}"
    return s


def embed_texts(sess, texts):
    """调 New-API /v1/embeddings；整批重试 2 次后逐条降级。返回 (向量列表, 失败索引集)。"""
    def call(batch):
        r = sess.post(f"{BASE}/v1/embeddings", json={
            "model": MODEL, "input": batch, "dimensions": DIM}, timeout=120)
        r.raise_for_status()
        data = r.json()["data"]
        return [d["embedding"] for d in sorted(data, key=lambda d: d["index"])]

    for attempt in range(3):
        try:
            return call(texts), set()
        except Exception as e:
            if attempt < 2:
                time.sleep(2 ** attempt)
                continue
            # 整批三败 → 逐条降级，坏行记错不断批
            vecs, bad = [], set()
            for i, t in enumerate(texts):
                try:
                    vecs.append(call([t])[0])
                except Exception as ee:
                    vecs.append(None)
                    bad.add(i)
                    click.echo(f"  [失败] 第{i}条: {ee}", err=True)
            return vecs, bad


def build_embed_text(payload: dict) -> str:
    """嵌入文本构造：payload 自带 embed_text 优先；否则固定拼接（与检索端语义对齐）。"""
    if payload.get("embed_text"):
        return payload["embed_text"]
    # 型取值补 type 键：升格卡 payload 用 type 存型（非 型/target_type），漏认会产出「【】名称…」丢型文本，
    # 令升格卡向量与检索端跨型语义错位；补一段式回退（additive，不动 型/target_type 既有行为）。
    t = payload.get("型") or payload.get("type") or payload.get("target_type", "")
    name = payload.get("名称", "")
    brief = payload.get("一句话摘要", "")
    fields = payload.get("字段") or {}
    body = "\n".join(f"{k}：{v}" for k, v in fields.items() if v and k not in ("名称", "一句话摘要"))
    return f"【{t}】{name}：{brief}\n{body}"[:4000]


class EmbeddingOwnershipConflict(RuntimeError):
    """同 hash 唯一行已归实体或其他活跃 draft，禁止迁移 owner。"""


def _embedding_owner_action(conn, draft_id, content_hash, *, lock=False):
    """判断同 hash 唯一行应幂等跳过还是写入；写段可锁行封住预查后的竞态。"""

    lock_clause = " FOR UPDATE OF e" if lock else ""
    owner = conn.execute(
        """SELECT e.draft_id, e.entity_id, e.deleted,
                  COALESCE(d.deleted, TRUE), d.tenant_id
           FROM example_knowledge_embedding e
           LEFT JOIN muse_knowledge_draft d ON d.id=e.draft_id
           WHERE e.tenant_id=%s AND e.content_hash=%s AND e.model=%s""" + lock_clause,
        (TENANT, content_hash, MODEL),
    ).fetchone()
    if not owner:
        return "write"

    owner_draft_id, owner_entity_id, embedding_deleted, owner_deleted, owner_tenant = owner
    # entity owner 是确认后的正式归属，任何 draft 都不得把它降级抢回。
    if owner_entity_id is not None:
        raise EmbeddingOwnershipConflict(
            f"同 hash 唯一行已归 entity：hash={content_hash}，entity={owner_entity_id}，"
            f"candidate={draft_id}"
        )
    # 只有当前租户、当前 draft、两侧都 active 才是真正的幂等命中。
    if owner_draft_id == draft_id:
        if owner_tenant != TENANT:
            raise EmbeddingOwnershipConflict(
                f"同 hash 当前 owner 租户不匹配：hash={content_hash}，"
                f"owner_tenant={owner_tenant}，candidate_tenant={TENANT}"
            )
        if not embedding_deleted and not owner_deleted:
            return "skip"
        return "write"
    # 空 owner、owner 行缺失或 owner draft 已软删时，可由当前活跃 draft 接管唯一行。
    if owner_draft_id is None or owner_deleted:
        return "write"
    raise EmbeddingOwnershipConflict(
        f"同 hash 唯一行已归其他 active draft：hash={content_hash}，"
        f"owner={owner_draft_id}，candidate={draft_id}"
    )


def _write_embedding(conn, draft_id, content_hash, text, vector):
    """在调用方单 draft 事务内锁定活性与 owner，条件写入并校验最终归属。"""

    # 写事务先按固定表顺序取得 ROW EXCLUSIVE 锁，避免与 reset 的多表锁形成交叉等待。
    conn.execute(
        "LOCK TABLE muse_knowledge_draft, example_knowledge_embedding IN ROW EXCLUSIVE MODE"
    )
    # 取得表锁后再锁 candidate draft：embed 先到时 reset 的七表 SHARE ROW EXCLUSIVE 会等待；
    # reset 先到时本查询等待其提交，随后读取 deleted=TRUE 并拒绝陈旧写入。
    candidate = conn.execute(
        """SELECT tenant_id, deleted, status, draft_payload FROM muse_knowledge_draft
           WHERE id=%s FOR UPDATE""",
        (draft_id,),
    ).fetchone()
    if not candidate:
        click.echo(f"  [跳过] draft={draft_id} 写前已不存在，未写向量", err=True)
        return False
    candidate_tenant, candidate_deleted, candidate_status, current_payload = candidate
    if candidate_tenant != TENANT:
        raise EmbeddingOwnershipConflict(
            f"draft 租户不匹配：draft={draft_id}，tenant={candidate_tenant}，expected={TENANT}"
        )
    if candidate_deleted:
        click.echo(f"  [跳过] draft={draft_id} 写前已软删，未写向量", err=True)
        return False
    if candidate_status != "pending":
        click.echo(
            f"  [跳过] draft={draft_id} 写前 status={candidate_status}，非 pending，未写向量",
            err=True,
        )
        return False

    # HTTP 期间 payload 可能被 parse/confirm 更新；锁内必须按当前 payload 重构文本与 hash，
    # 只要与 HTTP 请求所依据的快照不同，就丢弃陈旧向量，绝不覆盖并发产生的新结果。
    current_text = build_embed_text(current_payload or {})
    current_hash = hashlib.sha256(f"{current_text}|{MODEL}".encode()).hexdigest()
    if current_text != text or current_hash != content_hash:
        click.echo(
            f"  [跳过] draft={draft_id} 写前 payload/hash 漂移，"
            f"expected_hash={content_hash} current_hash={current_hash}，未写向量",
            err=True,
        )
        return False

    # 锁定该 draft 的全部活向量，保证 entity 归属和“每 draft 唯一活向量”在同一事务内判定。
    live_embeddings = conn.execute(
        """SELECT id, content_hash, entity_id FROM example_knowledge_embedding
           WHERE tenant_id=%s AND draft_id=%s AND deleted=FALSE
           FOR UPDATE""",
        (TENANT, draft_id),
    ).fetchall()
    entity_rows = [
        (row_id, row_hash, entity_id)
        for row_id, row_hash, entity_id in live_embeddings
        if entity_id is not None
    ]
    if entity_rows:
        raise EmbeddingOwnershipConflict(
            f"draft={draft_id} 存在 entity_id 非空旧活向量，禁止覆盖：{entity_rows}"
        )
    if any(row_hash == content_hash for _, row_hash, _ in live_embeddings):
        click.echo(f"  [跳过] draft={draft_id} 同 hash 活向量已由当前 draft 持有")
        return False

    action = _embedding_owner_action(conn, draft_id, content_hash, lock=True)
    if action == "skip":
        click.echo(f"  [跳过] draft={draft_id} 同 hash 活向量已由当前 draft 持有")
        return False

    if live_embeddings:
        # 当前 payload 已通过锁内 hash 重验，因此其余 hash 均为该 draft 的过期向量；
        # 只允许软删 draft owner，entity owner 已在上方失败关闭。
        conn.execute(
            """UPDATE example_knowledge_embedding SET deleted=TRUE, updater=%s
               WHERE tenant_id=%s AND draft_id=%s AND deleted=FALSE
                 AND entity_id IS NULL AND content_hash!=%s""",
            (ACTOR, TENANT, draft_id, content_hash),
        )

    # 条件 UPSERT 是行锁检查后的第二道防线：当预查时唯一行尚不存在、随后被并发插入时，
    # 仅允许当前 owner 或已失活 owner 迁移；entity/其他 active draft 均令 RETURNING 为空。
    upserted = conn.execute(
        """INSERT INTO example_knowledge_embedding
               (draft_id, content_hash, embed_text, model, dimensions, embedding,
                creator, updater, tenant_id)
           VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s)
           ON CONFLICT (tenant_id, content_hash, model)
           DO UPDATE SET draft_id=EXCLUDED.draft_id,
                         embed_text=EXCLUDED.embed_text,
                         model=EXCLUDED.model,
                         dimensions=EXCLUDED.dimensions,
                         embedding=EXCLUDED.embedding,
                         deleted=FALSE,
                         updater=EXCLUDED.updater
           WHERE example_knowledge_embedding.entity_id IS NULL
             AND (example_knowledge_embedding.draft_id=EXCLUDED.draft_id
                  OR NOT EXISTS (
                      SELECT 1 FROM muse_knowledge_draft owner
                      WHERE owner.id=example_knowledge_embedding.draft_id
                        AND owner.deleted=FALSE))
           RETURNING draft_id""",
        (draft_id, content_hash, text, MODEL, DIM, json.dumps(vector),
         ACTOR, ACTOR, TENANT),
    ).fetchone()
    if not upserted or upserted[0] != draft_id:
        raise EmbeddingOwnershipConflict(
            f"同 hash 唯一行未绑定当前 draft：hash={content_hash}，candidate={draft_id}"
        )
    return True


@click.command()
@click.option("--work-id", type=int, help="限定拆书批次的 work（draft.work_id=0 为全局行，用 source_id 关联参考书）")
@click.option("--limit", type=int, default=0, help="最多处理条数（0=不限）")
@click.option("--probe", help="自由文本试嵌（打印维度与前 5 维，不落库）")
def main(work_id, limit, probe):
    sess = _session()
    if probe:
        vecs, bad = embed_texts(sess, [probe])
        if bad:
            raise click.ClickException("试嵌失败")
        v = vecs[0]
        click.echo(f"维度={len(v)} 前5维={[round(x, 4) for x in v[:5]]}")
        return

    with psycopg.connect(DSN) as conn:
        # 待嵌=pending 草稿且无嵌入行
        sql = """SELECT d.id, d.draft_payload FROM muse_knowledge_draft d
                 WHERE d.tenant_id=%s AND d.deleted=FALSE AND d.status='pending'
                   AND NOT EXISTS (SELECT 1 FROM example_knowledge_embedding e
                                   WHERE e.tenant_id=%s AND e.draft_id=d.id AND e.deleted=FALSE)"""
        args = [TENANT, TENANT]
        if work_id is not None:
            sql += " AND d.source_id=%s"
            args.append(work_id)
        sql += " ORDER BY d.id"
        if limit:
            sql += f" LIMIT {int(limit)}"
        rows = conn.execute(sql, args).fetchall()
        click.echo(f"待嵌草稿: {len(rows)} 条")

        done = skip = fail = 0
        for i in range(0, len(rows), BATCH):
            chunk = rows[i:i + BATCH]
            texts, metas = [], []
            for did, payload in chunk:
                text = build_embed_text(payload or {})
                h = hashlib.sha256(f"{text}|{MODEL}".encode()).hexdigest()
                if _embedding_owner_action(conn, did, h) == "skip":
                    skip += 1  # 仅当前 active draft 持有 active 唯一行时才幂等跳过。
                    continue
                texts.append(text)
                metas.append((did, h, text))
            # owner 预查只用于避免无效 HTTP；明确提交读事务，HTTP 期间不持数据库事务或表锁。
            conn.commit()
            if not texts:
                continue
            vecs, bad = embed_texts(sess, texts)
            for j, (did, h, text) in enumerate(metas):
                if j in bad:
                    fail += 1
                    continue
                # 每个 draft 独立事务：draft 行锁、owner 行锁、条件 UPSERT 同生共死。
                with conn.transaction():
                    written = _write_embedding(conn, did, h, text, vecs[j])
                if written:
                    done += 1
                else:
                    skip += 1
            click.echo(f"  进度 {min(i + BATCH, len(rows))}/{len(rows)}（新嵌{done} 跳过{skip} 失败{fail}）")
        click.echo(f"完成：新嵌 {done}、跳过 {skip}、失败 {fail}")


if __name__ == "__main__":
    try:
        main()
    except (EmbeddingOwnershipConflict, psycopg.Error, requests.RequestException) as e:
        click.echo(f"[错误] {type(e).__name__}: {e}", err=True)
        sys.exit(1)