#!/usr/bin/env python3
"""
Test script to verify Mermaid diagram generation from LinkML schema.

This script:
1. Generates a Mermaid diagram from the schema
2. Verifies that hub connections are present
3. Validates the diagram structure

Usage:
    python test_mermaid_generation.py
"""

import sys
from pathlib import Path
from linkml_runtime.utils.schemaview import SchemaView

# Add scripts directory to path
SCRIPTS_DIR = Path(__file__).parent.parent / "scripts"
sys.path.insert(0, str(SCRIPTS_DIR))

from generate_mermaid_modular import generate_mermaid_from_schemaview

def test_hub_connections():
    """Test that hub connections are properly generated."""
    schema_path = Path(__file__).parent.parent / "schemas/20251121/linkml/01_custodian_name_modular.yaml"
    
    print("=" * 70)
    print("Testing Mermaid Diagram Generation from LinkML Schema")
    print("=" * 70)
    print()
    
    # Load schema
    print(f"📖 Loading schema: {schema_path.name}")
    sv = SchemaView(str(schema_path))
    print(f"✅ Loaded schema: {sv.schema.name}")
    print(f"   Classes: {len(list(sv.all_classes()))}")
    print(f"   Enums: {len(list(sv.all_enums()))}")
    print()
    
    # Test 1: Verify refers_to_custodian slot has correct range
    print("TEST 1: Base Slot Definition")
    print("-" * 70)
    base_slot = sv.get_slot("refers_to_custodian")
    print(f"  Slot: refers_to_custodian")
    print(f"    Range: {base_slot.range}")
    print(f"    Required: {base_slot.required}")
    
    assert base_slot.range == "Custodian", f"❌ FAIL: Expected range 'Custodian', got '{base_slot.range}'"
    assert base_slot.required == True, f"❌ FAIL: Expected required=True, got {base_slot.required}"
    print("✅ PASS: Base slot correctly defined")
    print()
    
    # Test 2: Verify induced slots in each class
    print("TEST 2: Induced Slot Definitions")
    print("-" * 70)
    
    test_classes = ['CustodianObservation', 'CustodianName', 'CustodianReconstruction']
    for class_name in test_classes:
        induced_slot = sv.induced_slot("refers_to_custodian", class_name)
        print(f"  Class: {class_name}")
        print(f"    Range: {induced_slot.range}")
        print(f"    Required: {induced_slot.required}")
        
        assert induced_slot.range == "Custodian", \
            f"❌ FAIL: {class_name} slot range should be 'Custodian', got '{induced_slot.range}'"
        assert induced_slot.required == True, \
            f"❌ FAIL: {class_name} slot should be required, got {induced_slot.required}"
    
    print("✅ PASS: All induced slots correctly defined")
    print()
    
    # Test 3: Generate Mermaid diagram
    print("TEST 3: Mermaid Diagram Generation")
    print("-" * 70)
    mermaid = generate_mermaid_from_schemaview(sv)
    print(f"  Generated diagram: {len(mermaid)} bytes")
    print()
    
    # Test 4: Verify hub connections in diagram
    print("TEST 4: Hub Connections in Diagram")
    print("-" * 70)
    
    expected_connections = [
        'CustodianReconstruction ||--|| Custodian : "refers_to_custodian"',
        'CustodianName ||--|| Custodian : "refers_to_custodian"',
        'CustodianObservation ||--|| Custodian : "refers_to_custodian"'
    ]
    
    for connection in expected_connections:
        if connection in mermaid:
            print(f"  ✅ {connection}")
        else:
            print(f"  ❌ MISSING: {connection}")
            print(f"\nGenerated diagram:\n{mermaid}")
            raise AssertionError(f"Expected connection not found: {connection}")
    
    print()
    print("✅ PASS: All hub connections present in diagram")
    print()
    
    # Test 5: Verify Custodian class definition
    print("TEST 5: Custodian Class Definition")
    print("-" * 70)
    
    custodian_class_start = mermaid.find("Custodian {")
    assert custodian_class_start > 0, "❌ FAIL: Custodian class not found in diagram"
    
    custodian_section = mermaid[custodian_class_start:custodian_class_start+200]
    print(f"  Found Custodian class definition:")
    print(f"    {custodian_section[:custodian_section.find('}')+1]}")
    
    assert "hc_id" in custodian_section, "❌ FAIL: hc_id not found in Custodian class"
    print("✅ PASS: Custodian class properly defined with hc_id")
    print()
    
    # Summary
    print("=" * 70)
    print("🎉 ALL TESTS PASSED")
    print("=" * 70)
    print()
    print("Summary:")
    print("  ✅ Base slot defined correctly (range: Custodian, required: True)")
    print("  ✅ Induced slots in all 3 classes correct")
    print("  ✅ Mermaid diagram generated successfully")
    print("  ✅ All 3 hub connections present in diagram")
    print("  ✅ Custodian class properly defined")
    print()
    print("The diagram generation script is working correctly!")
    print("Run this test after any schema changes to verify hub connections.")
    print()
    
    return True

if __name__ == '__main__':
    try:
        test_hub_connections()
        sys.exit(0)
    except AssertionError as e:
        print(f"\n❌ TEST FAILED: {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ ERROR: {e}", file=sys.stderr)
        import traceback
        traceback.print_exc()
        sys.exit(1)