#!/usr/bin/env python3
"""
Store content to kimi_kb (Knowledge Base) - Manual only

Usage:
    python3 kb_store.py "Content text" --title "Title" --domain "Category" --tags "tag1,tag2"
    python3 kb_store.py "Content" --title "X" --url "https://example.com" --source "docs.site"
"""

import json
import os
import sys
import urllib.request
import uuid
from datetime import datetime
from pathlib import Path

QDRANT_URL = "http://10.0.0.40:6333"
COLLECTION = "kimi_kb"
OLLAMA_URL = "http://10.0.0.10:11434/v1"

def get_embedding(text):
    """Generate embedding using snowflake-arctic-embed2"""
    data = json.dumps({
        "model": "snowflake-arctic-embed2",
        "input": text[:8192]
    }).encode()
    
    req = urllib.request.Request(
        f"{OLLAMA_URL}/embeddings",
        data=data,
        headers={"Content-Type": "application/json"}
    )
    
    try:
        with urllib.request.urlopen(req, timeout=60) as response:
            result = json.loads(response.read().decode())
            return result["data"][0]["embedding"]
    except Exception as e:
        print(f"Error generating embedding: {e}", file=sys.stderr)
        return None

def store_to_kb(text, title=None, url=None, source=None, domain=None, 
                tags=None, content_type="document"):
    """Store content to kimi_kb collection"""
    
    embedding = get_embedding(text)
    if embedding is None:
        return False
    
    point_id = str(uuid.uuid4())
    
    payload = {
        "text": text,
        "title": title or "Untitled",
        "url": url or "",
        "source": source or "manual",
        "domain": domain or "general",
        "tags": tags or [],
        "content_type": content_type,
        "date": datetime.now().strftime("%Y-%m-%d"),
        "created_at": datetime.now().isoformat(),
        "access_count": 0
    }
    
    point = {
        "points": [{
            "id": point_id,
            "vector": embedding,
            "payload": payload
        }]
    }
    
    data = json.dumps(point).encode()
    req = urllib.request.Request(
        f"{QDRANT_URL}/collections/{COLLECTION}/points?wait=true",
        data=data,
        headers={"Content-Type": "application/json"},
        method="PUT"
    )
    
    try:
        with urllib.request.urlopen(req, timeout=30) as response:
            result = json.loads(response.read().decode())
            return result.get("status") == "ok"
    except Exception as e:
        print(f"Error storing to KB: {e}", file=sys.stderr)
        return False

def main():
    import argparse
    
    parser = argparse.ArgumentParser(description="Store content to kimi_kb")
    parser.add_argument("content", help="Content to store")
    parser.add_argument("--title", default=None, help="Title of the content")
    parser.add_argument("--url", default=None, help="Source URL if from web")
    parser.add_argument("--source", default=None, help="Source name (e.g., 'docs.openclaw.ai')")
    parser.add_argument("--domain", default="general", help="Domain/category (e.g., 'OpenClaw', 'Docker')")
    parser.add_argument("--tags", default=None, help="Comma-separated tags")
    parser.add_argument("--type", default="document", choices=["document", "web", "code", "note"],
                       help="Content type")
    
    args = parser.parse_args()
    
    tags = [t.strip() for t in args.tags.split(",")] if args.tags else []
    
    print(f"Storing to kimi_kb: {args.title or 'Untitled'}...")
    
    if store_to_kb(
        text=args.content,
        title=args.title,
        url=args.url,
        source=args.source,
        domain=args.domain,
        tags=tags,
        content_type=args.type
    ):
        print(f"✅ Stored to kimi_kb ({args.domain})")
    else:
        print("❌ Failed to store")
        sys.exit(1)

if __name__ == "__main__":
    main()