#!/usr/bin/env python3
"""紧急恢复记忆 - 使用 urllib 替代 requests"""

import urllib.request
import urllib.error
import json
import time

DASHSCOPE_API_KEY = "sk-4111c9dba5334510968f9ae72728944e"
QDRANT_URL = "http://localhost:6333"
COLLECTION = "mem0_v4_shared"

def generate_embedding(text):
    """生成 1024 维向量"""
    data = json.dumps({"model": "text-embedding-v4", "input": [text]}).encode('utf-8')
    req = urllib.request.Request(
        "https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings",
        data=data,
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {DASHSCOPE_API_KEY}"
        }
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        result = json.loads(resp.read().decode('utf-8'))
        return result["data"][0]["embedding"]

def write_memory(id, memory_text, user_id, agent_id, metadata=None):
    """写入单条记忆到 Qdrant"""
    print(f"  生成向量...", end=" ", flush=True)
    embedding = generate_embedding(memory_text)
    print(f"✅ {len(embedding)} 维")
    
    payload = {
        "memory": memory_text,
        "user_id": user_id,
        "agent_id": agent_id,
        "source": "manual_recovery",
        "recovered_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    }
    if metadata:
        payload.update(metadata)
    
    data = json.dumps({
        "points": [{
            "id": id,
            "vector": embedding,
            "payload": payload
        }]
    }).encode('utf-8')
    
    req = urllib.request.Request(
        f"{QDRANT_URL}/collections/{COLLECTION}/points",
        data=data,
        headers={"Content-Type": "application/json"},
        method="PUT"
    )
    
    with urllib.request.urlopen(req, timeout=10) as resp:
        result = json.loads(resp.read().decode('utf-8'))
        print(f"  ✅ Qdrant: {result.get('status', 'unknown')}")
        return result

def search_memories(query, limit=5):
    """检索记忆"""
    query_embedding = generate_embedding(query)
    
    data = json.dumps({
        "vector": query_embedding,
        "limit": limit,
        "with_payload": True
    }).encode('utf-8')
    
    req = urllib.request.Request(
        f"{QDRANT_URL}/collections/{COLLECTION}/points/search",
        data=data,
        headers={"Content-Type": "application/json"}
    )
    
    with urllib.request.urlopen(req, timeout=10) as resp:
        result = json.loads(resp.read().decode('utf-8'))
        return result.get("result", [])

if __name__ == "__main__":
    print("=== 紧急恢复：犇犇部署记忆 ===\n")
    
    memories = [
        {
            "id": 1001,
            "text": "犇犇部署进度 - 企业微信插件挂载中，待完成",
            "user_id": "wang 院长",
            "agent_id": "main",
            "metadata": {"task_type": "deployment", "status": "pending", "priority": "high"}
        },
        {
            "id": 1002,
            "text": "Collection 名称统一为 mem0_v4_shared - 已完成（2026-02-27）",
            "user_id": "wang 院长",
            "agent_id": "main",
            "metadata": {"task_type": "configuration", "status": "completed"}
        },
        {
            "id": 1003,
            "text": "安全审计误报分析 - 已完成复核（2026-02-26）",
            "user_id": "wang 院长",
            "agent_id": "main",
            "metadata": {"task_type": "security", "status": "completed"}
        }
    ]
    
    for mem in memories:
        print(f"\n--- 记忆 {mem['id']} ---")
        write_memory(
            id=mem["id"],
            memory_text=mem["text"],
            user_id=mem["user_id"],
            agent_id=mem["agent_id"],
            metadata=mem["metadata"]
        )
        time.sleep(0.5)
    
    print("\n=== 验证检索 ===")
    results = search_memories("犇犇 部署", limit=5)
    print(f"检索到 {len(results)} 条记忆:\n")
    for r in results:
        print(f"  分数：{r['score']:.4f}")
        print(f"  内容：{r['payload'].get('memory', 'N/A')}")
        print()
    
    print("✅ 记忆恢复完成")