#!/usr/bin/env python3
"""
Enhanced memory storage with metadata support
Usage: store_memory.py "Memory text" [--tags tag1,tag2] [--importance medium] 
                         [--confidence high] [--source user|inferred|external]
                         [--verified] [--expires 2026-03-01] [--related id1,id2]
"""

import argparse
import json
import sys
import urllib.request
import uuid
from datetime import datetime, timedelta

QDRANT_URL = "http://10.0.0.40:6333"
COLLECTION_NAME = "kimi_memories"
OLLAMA_URL = "http://10.0.0.10:11434/v1"

def get_embedding(text):
    """Generate embedding using snowflake-arctic-embed2 via Ollama"""
    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=30) 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_memory(text, embedding, tags=None, importance="medium", date=None, 
                 source="conversation", confidence="high", source_type="user",
                 verified=True, expires_at=None, related_memories=None):
    """Store memory in Qdrant with enhanced metadata"""
    
    if date is None:
        date = datetime.now().strftime("%Y-%m-%d")
    
    # Generate a UUID for the point ID
    point_id = str(uuid.uuid4())
    
    # Build payload with all metadata
    payload = {
        "text": text,
        "date": date,
        "tags": tags or [],
        "importance": importance,
        "source": source,
        "confidence": confidence,  # high/medium/low
        "source_type": source_type,  # user/inferred/external
        "verified": verified,  # bool
        "created_at": datetime.now().isoformat(),
        "access_count": 0,
        "last_accessed": datetime.now().isoformat()
    }
    
    # Optional metadata
    if expires_at:
        payload["expires_at"] = expires_at
    if related_memories:
        payload["related_memories"] = related_memories
    
    # Qdrant upsert format
    upsert_data = {
        "points": [
            {
                "id": point_id,
                "vector": embedding,
                "payload": payload
            }
        ]
    }
    
    req = urllib.request.Request(
        f"{QDRANT_URL}/collections/{COLLECTION_NAME}/points?wait=true",
        data=json.dumps(upsert_data).encode(),
        headers={"Content-Type": "application/json"},
        method="PUT"
    )
    
    try:
        with urllib.request.urlopen(req, timeout=10) as response:
            result = json.loads(response.read().decode())
            if result.get("status") == "ok":
                return point_id
            else:
                print(f"Qdrant response: {result}", file=sys.stderr)
                return None
    except urllib.error.HTTPError as e:
        error_body = e.read().decode()
        print(f"HTTP Error {e.code}: {error_body}", file=sys.stderr)
        return None
    except Exception as e:
        print(f"Error storing memory: {e}", file=sys.stderr)
        return None

def link_memories(point_id, related_ids):
    """Link this memory to related memories (bidirectional)"""
    # Update this memory to include related
    # Then update each related memory to include this one
    pass  # Implementation would update existing points

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Store a memory in Qdrant with metadata")
    parser.add_argument("text", help="Memory text to store")
    parser.add_argument("--tags", help="Comma-separated tags")
    parser.add_argument("--importance", default="medium", choices=["low", "medium", "high"])
    parser.add_argument("--date", help="Date in YYYY-MM-DD format")
    parser.add_argument("--source", default="conversation", help="Source of the memory")
    parser.add_argument("--confidence", default="high", choices=["high", "medium", "low"],
                       help="Confidence in this memory's accuracy")
    parser.add_argument("--source-type", default="user", choices=["user", "inferred", "external"],
                       help="How this memory was obtained")
    parser.add_argument("--verified", action="store_true", default=True,
                       help="Whether this memory has been verified")
    parser.add_argument("--expires", help="Expiration date YYYY-MM-DD (for temporary memories)")
    parser.add_argument("--related", help="Comma-separated related memory IDs")
    
    args = parser.parse_args()
    
    # Parse tags and related memories
    tags = [t.strip() for t in args.tags.split(",")] if args.tags else []
    related = [r.strip() for r in args.related.split(",")] if args.related else None
    
    print(f"Generating embedding...")
    embedding = get_embedding(args.text)
    
    if embedding is None:
        print("❌ Failed to generate embedding", file=sys.stderr)
        sys.exit(1)
    
    print(f"Storing memory (vector dim: {len(embedding)})...")
    point_id = store_memory(
        args.text, embedding, tags, args.importance, args.date, args.source,
        args.confidence, args.source_type, args.verified, args.expires, related
    )
    
    if point_id:
        print(f"✅ Memory stored successfully")
        print(f"   ID: {point_id}")
        print(f"   Tags: {tags}")
        print(f"   Importance: {args.importance}")
        print(f"   Confidence: {args.confidence}")
        print(f"   Source: {args.source_type}")
        if args.expires:
            print(f"   Expires: {args.expires}")
    else:
        print(f"❌ Failed to store memory", file=sys.stderr)
        sys.exit(1)