#!/usr/bin/env python3
"""
Generate Mermaid ER diagrams from modular LinkML schemas.

This script works around the bug in gen-erdiagram that fails to resolve
linkml:types imports in modular schemas by using SchemaView.

Usage:
    python3 generate_mermaid_modular.py schema.yaml [output.mmd]
"""
import sys
from pathlib import Path
from linkml_runtime.utils.schemaview import SchemaView

def generate_mermaid_from_schemaview(sv: SchemaView) -> str:
    """
    Generate Mermaid ER diagram from SchemaView.
    
    This manually constructs the Mermaid syntax instead of using
    the buggy ErdiagramGenerator class.
    """
    lines = ["```mermaid"]
    lines.append("erDiagram")
    
    # Generate entities
    for class_name in sv.all_classes():
        cls = sv.get_class(class_name)
        
        lines.append(f"{class_name} {{")
        
        # Add attributes
        for slot_name in sv.class_slots(class_name):
            # Get slot definition (slot_usage or global)
            slot = None
            if cls.slot_usage and slot_name in cls.slot_usage:
                slot = cls.slot_usage[slot_name]
            else:
                slot = sv.get_slot(slot_name)
            
            if slot:
                slot_range = slot.range if slot.range else "string"
                
                # Skip complex types in attribute list (will be shown as relationships)
                if slot_range in sv.all_classes():
                    continue
                
                # Format: type attribute_name
                multivalued_marker = "List" if slot.multivalued else ""
                lines.append(f"    {slot_range}{multivalued_marker} {slot_name}  ")
        
        lines.append("}")
    
    lines.append("")
    
    # Generate relationships
    for class_name in sv.all_classes():
        cls = sv.get_class(class_name)
        
        # Inheritance relationships
        if cls.is_a:
            # Mermaid doesn't have inheritance arrows in ER diagrams
            # We can document this as a comment or skip
            pass
        
        # Association relationships
        for slot_name in sv.class_slots(class_name):
            # Use induced_slot to properly merge base slot with slot_usage
            slot = sv.induced_slot(slot_name, class_name)
            
            if slot and slot.range and slot.range in sv.all_classes():
                # Determine cardinality
                if slot.multivalued:
                    if slot.required:
                        # One-to-many required
                        cardinality = "||--}|"
                    else:
                        # One-to-many optional
                        cardinality = "||--}o"
                else:
                    if slot.required:
                        # One-to-one required
                        cardinality = "||--||"
                    else:
                        # One-to-one optional
                        cardinality = "||--|o"
                
                lines.append(f'{class_name} {cardinality} {slot.range} : "{slot_name}"')
    
    lines.append("")
    lines.append("```")
    lines.append("")
    
    return '\n'.join(lines)

def main():
    if len(sys.argv) < 2:
        print("Usage: generate_mermaid_modular.py <schema.yaml> [output.mmd]")
        print("\nGenerates Mermaid ER diagrams from modular LinkML schemas.")
        print("Works around gen-erdiagram bug with linkml:types imports.")
        sys.exit(1)
    
    schema_path = Path(sys.argv[1])
    output_path = Path(sys.argv[2]) if len(sys.argv) > 2 else None
    
    if not schema_path.exists():
        print(f"❌ Error: Schema file not found: {schema_path}", file=sys.stderr)
        sys.exit(1)
    
    # Load schema with SchemaView (handles modular imports)
    print(f"Loading schema: {schema_path}", file=sys.stderr)
    sv = SchemaView(str(schema_path))
    print(f"✅ Loaded schema: {sv.schema.name}", file=sys.stderr)
    print(f"   Classes: {len(list(sv.all_classes()))}", file=sys.stderr)
    print(f"   Enums: {len(list(sv.all_enums()))}", file=sys.stderr)
    
    # Generate Mermaid
    print("Generating Mermaid ER diagram...", file=sys.stderr)
    mermaid = generate_mermaid_from_schemaview(sv)
    
    # Output
    if output_path:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(mermaid)
        print(f"✅ Generated: {output_path}", file=sys.stderr)
        print(f"   Size: {len(mermaid)} bytes", file=sys.stderr)
    else:
        print(mermaid)

if __name__ == '__main__':
    main()