#!/usr/bin/env python3
"""
Full Pipeline Test for Memory Consolidation System
Tests all phases: Storage → Consolidation → Fact Extraction → Recall
"""

import requests
import time
import json

BASE_URL = "http://localhost:1865"

def send_message(text):
    """Send a message to Miku and get response"""
    resp = requests.post(f"{BASE_URL}/message", json={"text": text})
    return resp.json()

def get_qdrant_count(collection):
    """Get count of items in Qdrant collection"""
    resp = requests.post(
        f"http://localhost:6333/collections/{collection}/points/scroll",
        json={"limit": 1000, "with_payload": False, "with_vector": False}
    )
    return len(resp.json()["result"]["points"])

print("=" * 70)
print("🧪 FULL PIPELINE TEST - Memory Consolidation System")
print("=" * 70)

# TEST 1: Trivial Message Filtering
print("\n📋 TEST 1: Trivial Message Filtering")
print("-" * 70)

trivial_messages = ["lol", "k", "ok", "haha", "xd"]
important_message = "My name is Alex and I live in Seattle"

print("Sending trivial messages (should be filtered out)...")
for msg in trivial_messages:
    send_message(msg)
    time.sleep(0.5)

print("Sending important message...")
send_message(important_message)
time.sleep(1)

episodic_count = get_qdrant_count("episodic")
print(f"\n✅ Episodic memories stored: {episodic_count}")
if episodic_count < len(trivial_messages):
    print("   ✓ Trivial filtering working! (some messages were filtered)")
else:
    print("   ⚠️  Trivial filtering may not be active")

# TEST 2: Miku's Response Storage
print("\n📋 TEST 2: Miku's Response Storage")
print("-" * 70)

print("Sending message and checking if Miku's response is stored...")
resp = send_message("Tell me a very short fact about music")
miku_said = resp["content"]
print(f"Miku said: {miku_said[:80]}...")
time.sleep(2)

# Check for Miku's messages in episodic
resp = requests.post(
    "http://localhost:6333/collections/episodic/points/scroll",
    json={
        "limit": 100,
        "with_payload": True,
        "with_vector": False,
        "filter": {"must": [{"key": "metadata.speaker", "match": {"value": "miku"}}]}
    }
)
miku_messages = resp.json()["result"]["points"]
print(f"\n✅ Miku's messages in memory: {len(miku_messages)}")
if miku_messages:
    print(f"   Example: {miku_messages[0]['payload']['page_content'][:60]}...")
    print("   ✓ Bidirectional memory working!")
else:
    print("   ⚠️  Miku's responses not being stored")

# TEST 3: Add Rich Personal Information
print("\n📋 TEST 3: Adding Personal Information")
print("-" * 70)

personal_info = [
    "My name is Sarah Chen",
    "I'm 28 years old",
    "I work as a data scientist at Google",
    "My favorite color is blue",
    "I love playing piano",
    "I'm allergic to peanuts",
    "I live in Tokyo, Japan",
    "My hobbies include photography and hiking"
]

print(f"Adding {len(personal_info)} messages with personal information...")
for info in personal_info:
    send_message(info)
    time.sleep(0.5)

episodic_after = get_qdrant_count("episodic")
print(f"\n✅ Total episodic memories: {episodic_after}")
print(f"   ({episodic_after - episodic_count} new memories added)")

# TEST 4: Memory Consolidation
print("\n📋 TEST 4: Memory Consolidation & Fact Extraction")
print("-" * 70)

print("Triggering consolidation...")
resp = send_message("consolidate now")
consolidation_result = resp["content"]
print(f"\n{consolidation_result}")

time.sleep(2)

# Check declarative facts
declarative_count = get_qdrant_count("declarative")
print(f"\n✅ Declarative facts extracted: {declarative_count}")

if declarative_count > 0:
    # Show sample facts
    resp = requests.post(
        "http://localhost:6333/collections/declarative/points/scroll",
        json={"limit": 5, "with_payload": True, "with_vector": False}
    )
    facts = resp.json()["result"]["points"]
    print("\nSample facts:")
    for i, fact in enumerate(facts[:5], 1):
        print(f"   {i}. {fact['payload']['page_content']}")

# TEST 5: Fact Recall
print("\n📋 TEST 5: Declarative Fact Recall")
print("-" * 70)

queries = [
    "What is my name?",
    "How old am I?",
    "Where do I work?",
    "What's my favorite color?",
    "What am I allergic to?"
]

print("Testing fact recall with queries...")
correct_recalls = 0
for query in queries:
    resp = send_message(query)
    answer = resp["content"]
    print(f"\n❓ {query}")
    print(f"💬 Miku: {answer[:150]}...")
    
    # Basic heuristic: check if answer contains likely keywords
    keywords = {
        "What is my name?": ["Sarah", "Chen"],
        "How old am I?": ["28"],
        "Where do I work?": ["Google", "data scientist"],
        "What's my favorite color?": ["blue"],
        "What am I allergic to?": ["peanut"]
    }
    
    if any(kw.lower() in answer.lower() for kw in keywords[query]):
        print("   ✓ Correct recall!")
        correct_recalls += 1
    else:
        print("   ⚠️  May not have recalled correctly")
    
    time.sleep(1)

print(f"\n✅ Fact recall accuracy: {correct_recalls}/{len(queries)} ({correct_recalls/len(queries)*100:.0f}%)")

# TEST 6: Conversation History Recall
print("\n📋 TEST 6: Conversation History (Episodic) Recall")
print("-" * 70)

print("Asking about conversation history...")
resp = send_message("What have we talked about today?")
summary = resp["content"]
print(f"💬 Miku's summary:\n{summary}")

# Final Summary
print("\n" + "=" * 70)
print("📊 FINAL SUMMARY")
print("=" * 70)
print(f"✅ Episodic memories: {get_qdrant_count('episodic')}")
print(f"✅ Declarative facts: {declarative_count}")
print(f"✅ Miku's messages stored: {len(miku_messages)}")
print(f"✅ Fact recall accuracy: {correct_recalls}/{len(queries)}")

# Overall verdict
if declarative_count >= 5 and correct_recalls >= 3:
    print("\n🎉 PIPELINE TEST: PASS")
    print("   All major components working correctly!")
else:
    print("\n⚠️  PIPELINE TEST: PARTIAL PASS")
    print("   Some components may need adjustment")

print("\n" + "=" * 70)