#!/usr/bin/env python3
"""
插入 AI 创作工作流的 16 个通用能力模块知识条目
"""
import os
import sys
import json
import uuid
from datetime import datetime, timezone
from pathlib import Path

# 添加项目路径
sys.path.insert(0, '/root/Agent')

os.environ['OPEN_ROUTER_API_KEY'] = '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'
os.environ['QWEN_API_KEY'] = 'sk-9453c827b9e14108b53d2b30ef7c75fe'
os.environ['KNOWHUB_DB'] = 'gp-t4n72471pkmt4b9q7o-master.gpdbmaster.singapore.rds.aliyuncs.com'
os.environ['KNOWHUB_PORT'] = '5432'
os.environ['KNOWHUB_USER'] = 'aiddit_aigc'
os.environ['KNOWHUB_PASSWORD'] = '%a&&yqNxg^V1$toJ*WOa^-b^X=QJ'
os.environ['KNOWHUB_DB_NAME'] = 'knowhub'

from knowhub.knowhub_db.pg_store import PostgreSQLStore
from knowhub.embeddings import get_embedding

# 知识数据 - AI 创作工作流的 16 个通用能力模块
knowledge_data = {
    "task": "AI 创作工作流的 16 个通用能力模块",
    "content": """## 能力模块清单

### 素材准备类
1. **准备参考素材**：收集或生成用于后续创作的参考图像、纹理、角色基础等素材
2. **生成参考图表**：生成角色参考表、多视角图、多表情图等参考资料

### 角色定义类
3. **定义角色特征**：设定人物面部特征、表情状态、身份一致性等角色属性

### 场景构建类
4. **构建场景环境**：创建或合成场景背景、虚拟环境、空间布局

### 视觉控制类
5. **配置光影效果**：设置光源方向、灯光类型、明暗对比等光照条件
6. **控制姿态构图**：控制人物姿态、画面构图、视角、摄影参数
7. **应用风格控制**：应用风格迁移、保持视觉一致性、注入特定美学风格
8. **处理纹理效果**：添加纹理、颗粒、质感、色调等视觉效果
9. **创建纹理布局**：创建四方连续、渐变网格、矢量图形等纹理布局

### 生成执行类
10. **执行内容生成**：执行核心的图像、视频、渲染生成操作
11. **配置控制技术**：配置 ControlNet、深度图、姿态图等控制条件
12. **配置工作流**：配置工具节点、构建流程、设置参数

### 后期处理类
13. **修复细节瑕疵**：对局部区域进行重绘、修复瑕疵、优化细节
14. **增强图像质量**：提升图像分辨率、质量增强、最终精修
15. **后期合成处理**：抠图、选区创建、图像合成、图层处理
16. **视频后期处理**：视频剪辑、颜色分级、特效增强

## 跨品类复用性
以上模块均从 5 个不同品类的工作流（角色创作、人像精修、肖像摄影、光影控制、纹理设计）中抽象得出，可跨品类复用。""",
    "types": ["strategy"],
    "tags": {
        "domain": "AI 创作",
        "category": "能力模块",
        "source": "工作流分析 pipeline"
    },
    "score": 5,
    "source": {"category": "research"}
}

print("正在连接数据库...")
store = PostgreSQLStore()

# 生成 ID
now = datetime.now(timezone.utc)
knowledge_id = f"knowledge-{now.strftime('%Y%m%d')}-{uuid.uuid4().hex[:8]}"
message_id = f"msg-{uuid.uuid4().hex[:12]}"

# 生成 embedding
print("正在生成 embedding...")
task_embedding = get_embedding(knowledge_data["task"])
content_embedding = get_embedding(knowledge_data["content"])

# 构建知识记录
knowledge = {
    "id": knowledge_id,
    "message_id": message_id,
    "task": knowledge_data["task"],
    "content": knowledge_data["content"],
    "types": knowledge_data["types"],
    "tags": knowledge_data["tags"],
    "tag_keys": list(knowledge_data["tags"].keys()),
    "scopes": [],
    "owner": "librarian",
    "source": knowledge_data["source"],
    "eval": {"score": knowledge_data["score"]},
    "created_at": now.isoformat(),
    "updated_at": now.isoformat(),
    "status": "approved",
    "task_embedding": task_embedding,
    "content_embedding": content_embedding,
    "requirement_ids": [],
    "capability_ids": [],
    "tool_ids": [],
    "resource_ids": [],
}

print(f"正在插入知识：{knowledge_id}")
print(f"  Task: {knowledge['task']}")
print(f"  Types: {knowledge['types']}")
print(f"  Tags: {knowledge['tags']}")
print(f"  Content 长度：{len(knowledge['content'])} 字符")

store.insert(knowledge)

print(f"\n✓ 知识已成功入库!")
print(f"  ID: {knowledge_id}")
print(f"  Message ID: {message_id}")

# 验证插入
retrieved = store.get_by_id(knowledge_id)
if retrieved:
    print(f"\n✓ 验证成功：已检索到刚插入的知识")
    print(f"  Task: {retrieved['task']}")
    print(f"  Content 长度：{len(retrieved['content'])} 字符")
    print(f"  Types: {retrieved['types']}")
    print(f"  Tags: {retrieved['tags']}")
else:
    print(f"\n✗ 验证失败：无法检索到刚插入的知识")

store.close()
print("\n数据库连接已关闭")