import cocoindex as coco
from cocoindex.connectors import localfs
from cocoindex.ops.text import RecursiveSplitter

# 定义文件处理函数
@coco.fn(memo=True)
async def process_file(file, table):
    text = await file.read_text()
    chunks = splitter.split(text, chunk_size=2000)
    for chunk in chunks:
        table.declare_row(
            text=chunk.text,
            embedding=embed(chunk.text)
        )

# 创建 App 并运行
app = coco.App("my_index", main, src="./docs")
app.update_blocking()

# 安装 CocoIndex
pip install -U cocoindex

# 验证安装
python -c "import cocoindex; print(cocoindex.__version__)"

# 创建项目目录
mkdir my-cocoindex-app && cd my-cocoindex-app

# 配置本地数据库（可选，默认使用 SQLite）
echo "COCOINDEX_DB=./cocoindex.db" > .env

from cocoindex.connectors import localfs
from cocoindex.resources.file import PatternFilePathMatcher

files = localfs.walk_dir(
    sourcedir,
    recursive=True,
    path_matcher=PatternFilePathMatcher(
        included_patterns=["**/*.md"],
    ),
    live=True,
)

# files.items() 返回 (StableKey, FileLike) 对
# 每个 FileLike 可以 read_text() 读取内容

@coco.fn(memo=True)
async def process_file(
    file: FileLike,
    table: postgres.TableTarget[DocEmbedding],
) -> None:
    # 1. 读取文件内容
    text = await file.read_text()
    
    # 2. 文本分块
    chunks = _splitter.split(
        text,
        chunk_size=2000,
        chunk_overlap=500,
        language="markdown",
    )
    
    # 3. 为每个 chunk 生成 embedding 并声明目标行
    id_gen = IdGenerator()
    await coco.map(
        process_chunk, chunks,
        file.file_path.path, id_gen, table
    )

@coco.fn
async def app_main(sourcedir: pathlib.Path) -> None:
    # 1. 挂载目标表
    target_table = await postgres.mount_table_target(
        PG_DB,
        table_name="doc_embeddings",
        table_schema=schema,
    )
    target_table.declare_vector_index(
        column="embedding"
    )
    
    # 2. 扫描源文件
    files = localfs.walk_dir(sourcedir, recursive=True)
    
    # 3. 为每个文件创建独立处理组件
    await coco.mount_each(
        process_file,
        files.items(),
        target_table
    )

# 定义上下文键
PG_DB = coco.ContextKey[asyncpg.Pool]("db")
EMBEDDER = coco.ContextKey[SentenceTransformerEmbedder](
    "embedder",
    detect_change=True  # 模型变化时触发重新计算
)

# 定义生命周期——创建和清理共享资源
@coco.lifespan
async def coco_lifespan(builder):
    async with asyncpg.create_pool(DATABASE_URL) as pool:
        builder.provide(PG_DB, pool)
        builder.provide(EMBEDDER, SentenceTransformerEmbedder(model))
        yield  # 生命周期持续到 App 停止

# 在变换函数中使用上下文
embedding = await coco.use_context(EMBEDDER).embed(text)

from dataclasses import dataclass
from typing import Annotated
from numpy.typing import NDArray

# 定义目标表的数据结构
@dataclass
class DocEmbedding:
    id: int
    filename: str
    text: str
    embedding: Annotated[NDArray, EMBEDDER]
    chunk_start: int
    chunk_end: int

# Annotated[NDArray, EMBEDDER] 告诉 CocoIndex：
# embedding 字段由 EMBEDDER ContextKey 自动生成
# 当模型变化时 (detect_change=True)，自动重新计算

from cocoindex.connectors import postgres

# 挂载目标表
target_table = await postgres.mount_table_target(
    PG_DB,  # ContextKey 绑定的连接池
    table_name="doc_embeddings",
    table_schema=await postgres.TableSchema.from_class(
        DocEmbedding,
        primary_key=["id"],
    ),
    pg_schema_name="coco_examples",
)

# 声明向量索引
target_table.declare_vector_index(
    column="embedding"
)

from cocoindex.ops.text import RecursiveSplitter, detect_code_language

@coco.fn(memo=True)
async def process_file(file: FileLike, table):
    text = await file.read_text()
    
    # 自动检测编程语言
    language = detect_code_language(
        filename=str(file.file_path.path.name)
    )
    
    # 语法感知分块——不会在函数中间切断
    chunks = _splitter.split(
        text,
        chunk_size=1000,
        min_chunk_size=300,
        chunk_overlap=300,
        language=language,  # Python → AST 感知
    )
    
    id_gen = IdGenerator()
    await coco.map(process_chunk, chunks,
              file.file_path.path, id_gen, table)

async def query_once(pool, embedder, query, *, top_k=5):
    # 1. 将查询文本转换为向量
    query_vec = await embedder.embed(query)
    
    # 2. 在 Postgres 中搜索最相似的向量
    async with pool.acquire() as conn:
        rows = await conn.fetch(
            f"""
            SELECT filename, code,
                   embedding <=> $1 AS distance,
                   start_line, end_line
            FROM "{SCHEMA}"."{TABLE}"
            ORDER BY distance ASC
            LIMIT $2
            """,
            query_vec,
            top_k,
        )
    
    # 3. 打印结果
    for r in rows:
        score = 1.0 - float(r["distance"])
        print(f"[{score:.3f}] {r['filename']}")

# 运行查询（命令行传入查询文本）
$ python main.py "how to split text into chunks"

# 输出示例：
[0.847] python/cocoindex/ops/text.py (L120-L155)
    def split(self, text, chunk_size, ...):
        raw_chunks = self._splitter.split(
            text, chunk_size, min_chunk_size, ...
        )
---
[0.792] python/cocoindex/ops/text.py (L39-L84)
    class SeparatorSplitter:
        def split(self, text) -> list[Chunk]:
---

# 也可以进入 REPL 模式
$ python main.py
Enter search query (or Enter to quit): embedding model configuration
[0.831] python/cocoindex/ops/sentence_transformers.py

# 命令行方式：-L 标志启用 Live 模式
$ cocoindex update -L main

# 或者代码方式：
app = coco.App("CodeEmbedding", app_main, sourcedir="./src")

# 同步运行一次（扫描 + 处理 + 完成）
app.update_blocking()

# Live 模式（持续监控，发现变化自动更新）
handle = app.update(live=True)

# 监控进度
async for snapshot in handle.watch():
    print(f"Status: {snapshot.status}")
    print(f"Stats: {snapshot.stats}")

# 获取最终结果
result = await handle.result()

# 查看帮助
$ cocoindex --help

# 运行更新（一次性）
$ cocoindex update main

# 运行更新（Live 模式，持续监控）
$ cocoindex update -L main

# 显示进度详情
$ cocoindex update --report main

# 删除应用（清除所有目标状态）
$ cocoindex drop main

# 列出所有应用
$ cocoindex ls

# 全量重处理（忽略缓存）
$ cocoindex update --full-reprocess main

# 创建 App
app = coco.App(
    coco.AppConfig(
        name="CodeEmbedding",
        max_inflight_components=1024,
    ),
    app_main,
    sourcedir=pathlib.Path("./src"),
)

# 模式 1：同步阻塞（适合脚本、CLI）
result = app.update_blocking(
    report_to_stdout=True,
)

# 模式 2：异步非阻塞（适合 Web 服务）
handle = app.update()
result = await handle.result()

# 模式 3：Live 模式（持续监控）
handle = app.update(live=True)
async for snap in handle.watch():
    if snap.status == UpdateStatus.READY:
        print("Initial processing complete!")

# 注册异常处理器
@coco.exception_handler
async def handle_error(exc, ctx):
    logger.error(
        f"Component {ctx.stable_path} failed: "
        f"{ctx.processor_name} in {ctx.env_name}"
    )

# 监控更新进度
handle = app.update(live=True)
async for snapshot in handle.watch():
    for name, stats in snapshot.stats.by_component.items():
        print(f"[{name}] "
              f"processed={stats.num_processed} "
              f"cached={stats.num_cached}")

# 1. 将 CocoIndex skill 添加到项目
# 复制 skills/cocoindex/ 到 .claude/skills/
cp -r skills/cocoindex/ .claude/skills/

# 2. Claude Code 自动识别 CocoIndex API
# 正确使用 @coco.fn、mount_each、ContextKey 等

# 3. 管道代码示例
@coco.fn(memo=True)
async def index_file(file, table):
    for chunk in RecursiveSplitter().split(
        await file.read_text()
    ):
        table.declare_row(
            text=chunk.text,
            embedding=embed(chunk.text)
        )