"""Comprehensive MONIKA system audit script."""

import sys
sys.path.insert(0, '.')

from salience_os_seed.conversation.session import ConversationSession, ConversationConfig
import json

def audit_system():
    """Run comprehensive system audit."""
    
    print("=" * 60)
    print("MONIKA SYSTEM AUDIT")
    print("=" * 60)
    
    # Load session
    print("\n1. Loading checkpoint...")
    config = ConversationConfig()
    config.lm.checkpoint_path = 'storage/proto_lm/dolly15k.pt'
    config.learning_enabled = True
    
    session = ConversationSession(config)
    runtime = session.runtime
    proto_lm = session.proto_lm
    
    print(f"   [OK] Loaded: step={proto_lm.step}, vocab={proto_lm.vocab.size()}")
    
    # Check model parameters
    print("\n2. Model architecture...")
    try:
        param_count = sum(p.numel() for p in proto_lm.model.parameters())
        print(f"   [OK] Total parameters: {param_count:,}")
    except Exception as e:
        print(f"   [FAIL] Failed to count parameters: {e}")
    
    # Test training
    print("\n3. Testing training step...")
    try:
        test_text = "The quick brown fox jumps over the lazy dog."
        old_step = proto_lm.step
        proto_lm.training_step(test_text)
        new_step = proto_lm.step
        print(f"   [OK] Training works: step {old_step} → {new_step}")
    except Exception as e:
        print(f"   [FAIL] Training failed: {e}")
    
    # Test generation
    print("\n4. Testing text generation...")
    try:
        prompt = "Hello world"
        generated = proto_lm.sample(prompt, max_tokens=30)
        print(f"   Prompt: '{prompt}'")
        print(f"   Generated: '{generated}'")
        if generated == prompt or generated.strip() == prompt.strip():
            print(f"   [FAIL] Generation broken: just echoing prompt")
        else:
            print(f"   [OK] Generation works")
    except Exception as e:
        print(f"   [FAIL] Generation failed: {e}")
    
    # Test runtime step
    print("\n5. Testing runtime step...")
    try:
        metrics = session.process_user_input("Test message")
        print(f"   [OK] Runtime step works")
        print(f"     Step: {metrics.step}")
        print(f"     Budget left: {metrics.budget_left:.1f}")
        print(f"     Verification: {metrics.verification_passed}")
    except Exception as e:
        print(f"   [FAIL] Runtime step failed: {e}")
    
    # Check yearning state
    print("\n6. Checking yearning state...")
    try:
        yearning = runtime.introspection.get_yearning_state(refresh=True)
        print(f"   [OK] Yearning state retrieved: {len(yearning)} entries")
        
        # Check if yearning is computed
        sample_key = list(yearning.keys())[0]
        sample = yearning[sample_key]
        
        has_yearning_field = 'yearning' in sample
        print(f"   Has 'yearning' field: {has_yearning_field}")
        print(f"   Fields present: {list(sample.keys())}")
        
        # Manually compute yearning
        computed = sample['desire'] * (1 - sample['saturation'])
        print(f"   Manual yearning calc: {computed:.3f}")
        
        if not has_yearning_field:
            print(f"   [FAIL] BUG: yearning field missing from dict")
        
    except Exception as e:
        print(f"   [FAIL] Yearning state failed: {e}")
    
    # Check controller dynamics
    print("\n7. Checking controller dynamics...")
    try:
        dynamics = runtime.introspection.controller_snapshot()
        print(f"   [OK] Controller snapshot works")
        print(f"   Last action: {dynamics.get('last_action')}")
        if 'scores' in dynamics:
            print(f"   Action scores available: {len(dynamics['scores'])} actions")
    except Exception as e:
        print(f"   [FAIL] Controller snapshot failed: {e}")
    
    # Check memory
    print("\n8. Checking structured memory...")
    try:
        memory_snapshot = runtime.memory.as_runtime_mapping()
        facts = len(memory_snapshot.get('facts', []))
        todos = len(memory_snapshot.get('todos', []))
        hypotheses = len(memory_snapshot.get('hypotheses', []))
        print(f"   [OK] Memory accessible")
        print(f"     Facts: {facts}, Todos: {todos}, Hypotheses: {hypotheses}")
    except Exception as e:
        print(f"   [FAIL] Memory check failed: {e}")
    
    # Test memory operations
    print("\n9. Testing memory operations...")
    try:
        session.memory_operator.execute({
            'op': 'add_fact',
            'text': 'Test fact from audit',
            'score': 1.0
        })
        print(f"   [OK] Memory operations work")
    except Exception as e:
        print(f"   [FAIL] Memory operations failed: {e}")
    
    # Test response generation
    print("\n10. Testing full response generation...")
    try:
        snapshot = session.generate_response(prompt="Tell me about AI")
        print(f"   [OK] Response generation works")
        print(f"   Response: '{snapshot.response[:100]}...'")
    except Exception as e:
        print(f"   [FAIL] Response generation failed: {e}")
    
    print("\n" + "=" * 60)
    print("AUDIT COMPLETE")
    print("=" * 60)

if __name__ == '__main__':
    audit_system()