#!/usr/bin/env python3
"""Quick test of the knowledge graph module."""
import json
import sys
sys.path.insert(0, '.')
from graph import (
    add_entity, add_relation, get_entity, search_entities,
    explore, health, merge_entities, init_db, _rebuild_index,
    VALID_RELATION_TYPES
)

# Ensure schema is initialized
init_db()

# Clean up any test data from previous runs
from graph import DB_PATH
import sqlite3
conn = sqlite3.connect(DB_PATH)
conn.execute('DELETE FROM relations')
conn.execute('DELETE FROM entities')
conn.commit()
conn.close()

# Rebuild empty index
_rebuild_index()

# Test health
print('=== HEALTH ===')
h = health()
print(json.dumps(h, indent=2, default=str))

# Test adding entities
print('\n=== ADD ENTITIES ===')
e1 = add_entity('person', 'Matt Hoffmann',
                'Matt lives in Green Bay, WI. Husband to Aundrea, father to Sullivan and Harper.',
                tags=['green_bay', 'family'], confidence=1.0)
print(f'Entity 1: {e1}')

e2 = add_entity('person', 'Aundrea Hoffmann',
                'Aundrea is Matts wife.',
                tags=['family'])
print(f'Entity 2: {e2}')

e3 = add_entity('place', 'Green Bay',
                'Green Bay, Wisconsin. Home of the Hoffmann family.',
                tags=['wisconsin', 'home'])
print(f'Entity 3: {e3}')

# Test adding relations
print('\n=== ADD RELATIONS ===')
r1 = add_relation(e1, e2, 'spouse_of', confidence=1.0)
print(f'Relation 1: {r1}')
r2 = add_relation(e1, e3, 'lives_in', confidence=1.0)
print(f'Relation 2: {r2}')

# Test get_entity
print('\n=== GET ENTITY ===')
entity = get_entity(e1)
print(f'Entity: {entity["subject"]} ({entity["type"]})')
print(f'Relations out: {len(entity["relations_out"])}')
print(f'Relations in: {len(entity["relations_in"])}')
for rel in entity['relations_out']:
    print(f'  -> {rel["relation_type"]} -> {rel["target_subject"]}')

# Test explore
print('\n=== EXPLORE ===')
ex = explore(e1, depth=1)
print(f'Root: {ex["root"]["subject"]}')
for nid, node in ex.get('nodes', {}).items():
    print(f'  Neighbor: {node["subject"]} ({node["type"]})')

# Test semantic search
print('\n=== SEARCH ===')
results = search_entities('Matts family', top_k=3)
for r in results:
    print(f'  [{r["similarity"]:.3f}] {r["subject"]} ({r["type"]})')

# Test dedup
print('\n=== DEDUP ===')
e4 = add_entity('person', 'matt hoffmann',
                'Lives in WI with his family.',
                tags=['green_bay'])
print(f'Dedup returned: {e4} (should equal {e1})')
assert e4 == e1, 'Dedup should return existing entity ID'
print('  ✓ Dedup working')

# Test vague relation rejection
print('\n=== VAGUE RELATION REJECTION ===')
try:
    add_relation(e1, e3, 'related_to')
    print('  FAILED: Should have rejected "related_to"')
except ValueError as e:
    print(f'  ✓ Rejected: {e}')

print('\n=== All tests passed! ===')