"""知识切片与检索（M3）

Revision ID: 0006
Revises: 0005
Create Date: 2026-10-29

建 `chunk` 表：知识正文按段切开，每片一行，带向量数组。

两个刻意的取舍：

1. **向量存 `real[]`，不依赖 pgvector**。本机拿不到带 pgvector 的实例（Docker 起不来，
   fnsvr-01 明确不动），但检索链路不能一直停在纸面上。`real[]` + 一条 SQL 余弦函数
   在小数据量下够用；将来换上 pgvector 只需改列类型、加 HNSW 索引，应用层接口不变。
2. **检索范围只能是 `published`**。草稿是"还没被人认可的话"，进检索等于让 AI 把
   猜测当结论引用出去——这条过滤写在查询里，不靠调用方自觉。
"""

from __future__ import annotations

from collections.abc import Sequence

from alembic import op

revision: str = "0006"
down_revision: str = "0005"
branch_labels: Sequence[str] | None = None
depends_on: Sequence[str] | None = None

# 余弦相似度。写成 SQL 函数而不是在 Python 里算，是为了让过滤与排序留在数据库里做——
# 把整表向量捞回应用算相似度，数据一多就会拖垮 API。
COSINE_SQL = """
CREATE OR REPLACE FUNCTION app_cosine(a real[], b real[]) RETURNS double precision
LANGUAGE plpgsql IMMUTABLE PARALLEL SAFE AS $fn$
DECLARE
    n int;
    dot double precision := 0;
    na double precision := 0;
    nb double precision := 0;
BEGIN
    n := LEAST(array_length(a, 1), array_length(b, 1));
    IF n IS NULL THEN
        RETURN NULL;
    END IF;
    FOR i IN 1..n LOOP
        dot := dot + a[i] * b[i];
        na := na + a[i] * a[i];
        nb := nb + b[i] * b[i];
    END LOOP;
    IF na = 0 OR nb = 0 THEN
        RETURN 0;
    END IF;
    RETURN dot / sqrt(na * nb);
END
$fn$;
"""

CREATE_SQL = """
CREATE TABLE chunk (
    id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
    workspace_id uuid NOT NULL,
    item_id uuid NOT NULL REFERENCES item(id) ON DELETE CASCADE,
    idx integer NOT NULL,
    heading varchar(200),
    content text NOT NULL,
    char_count integer NOT NULL DEFAULT 0,
    embedding real[],
    embed_model varchar(64),
    version integer NOT NULL DEFAULT 1,
    created_at timestamptz NOT NULL DEFAULT now(),
    updated_at timestamptz NOT NULL DEFAULT now(),
    CONSTRAINT uq_chunk_item_idx UNIQUE (item_id, idx)
)
"""

# asyncpg 的执行器不接受多语句，一条 `execute` 只能带一句 SQL（0001 踩过同样的坑）
INDEX_STATEMENTS = (
    "CREATE INDEX ix_chunk_ws_item ON chunk (workspace_id, item_id)",
    "COMMENT ON TABLE chunk IS '知识正文切片，检索的最小单位'",
)

POLICY_SQL = (
    "ALTER TABLE chunk ENABLE ROW LEVEL SECURITY",
    "CREATE POLICY chunk_isolation ON chunk "
    "USING (workspace_id = app_ws_id()) WITH CHECK (workspace_id = app_ws_id())",
    f"GRANT SELECT, INSERT, UPDATE, DELETE ON chunk TO {APP_ROLE}",
)


def upgrade() -> None:
    op.execute(COSINE_SQL)
    op.execute(CREATE_SQL)
    for statement in INDEX_STATEMENTS:
        op.execute(statement)
    for statement in POLICY_SQL:
        op.execute(statement)


def downgrade() -> None:
    op.execute("DROP POLICY IF EXISTS chunk_isolation ON chunk")
    op.execute("DROP TABLE IF EXISTS chunk")
    # 函数被别的东西引用时先别删；本轮只有切片在用，留一条注释说明为什么还留着。
    # op.execute("DROP FUNCTION IF EXISTS app_cosine(real[], real[])")