#!/usr/bin/env python3
"""
Generate PlantUML class diagrams from modular LinkML schemas.

This script works around the bug in gen-plantuml that fails to resolve
linkml:types imports in modular schemas by using SchemaView instead of
SchemaLoader (same approach as gen-erdiagram).

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

def generate_plantuml_from_schemaview(sv: SchemaView) -> str:
    """
    Generate PlantUML class diagram from SchemaView.
    
    This manually constructs the PlantUML syntax instead of using
    the buggy PlantumlGenerator class.
    """
    lines = ["@startuml"]
    lines.append("!theme plain")
    lines.append("skinparam linetype ortho")
    lines.append("skinparam groupInheritance 2")
    lines.append("")
    
    # Add header comments
    lines.append(f"' Generated from LinkML schema: {sv.schema.name}")
    lines.append(f"' Schema ID: {sv.schema.id}")
    if sv.schema.description:
        lines.append(f"' Description: {sv.schema.description}")
    lines.append(f"' Generated by: generate_plantuml_modular.py (workaround for gen-plantuml bug)")
    lines.append("")
    lines.append("' ============================================")
    lines.append("' CLASSES")
    lines.append("' ============================================")
    lines.append("")
    
    # Generate classes
    for class_name in sv.all_classes():
        cls = sv.get_class(class_name)
        
        # Skip if abstract (we'll show inheritance)
        if cls.abstract:
            lines.append(f"abstract class {class_name} {{")
        else:
            lines.append(f"class {class_name} {{")
        
        # Add attributes
        # Use slot_usage (class-specific) + get_slot (global) to avoid induced_slot issues
        for slot_name in sv.class_slots(class_name):
            # Try class-specific slot_usage first (has correct ranges)
            slot = None
            if cls.slot_usage and slot_name in cls.slot_usage:
                slot = cls.slot_usage[slot_name]
            else:
                # Fall back to global slot definition
                slot = sv.get_slot(slot_name)
            
            if slot:
                slot_range = slot.range if slot.range else "string"
                required = "!" if slot.required else ""
                multivalued = "[*]" if slot.multivalued else ""
                lines.append(f"  +{slot_name}: {slot_range}{multivalued} {required}")
        
        lines.append("}")
        lines.append("")
    
    lines.append("' ============================================")
    lines.append("' ENUMERATIONS")
    lines.append("' ============================================")
    lines.append("")
    
    # Generate enums
    for enum_name in sv.all_enums():
        enum = sv.get_enum(enum_name)
        if enum.description:
            lines.append(f"' {enum.description}")
        lines.append(f"enum {enum_name} {{")
        if enum.permissible_values:
            for pv_name in enum.permissible_values:
                lines.append(f"  {pv_name}")
        lines.append("}")
        lines.append("")
    
    lines.append("' ============================================")
    lines.append("' INHERITANCE")
    lines.append("' ============================================")
    lines.append("")
    
    # Generate inheritance relationships
    for class_name in sv.all_classes():
        cls = sv.get_class(class_name)
        if cls.is_a:
            lines.append(f"{cls.is_a} <|-- {class_name}")
    
    lines.append("")
    
    lines.append("' ============================================")
    lines.append("' ASSOCIATIONS")
    lines.append("' ============================================")
    lines.append("")
    
    # Generate associations
    for class_name in sv.all_classes():
        cls = sv.get_class(class_name)
        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 and slot.range and slot.range in sv.all_classes():
                card = '"1..*"' if slot.multivalued else '"1"'
                req_card = '"1"' if slot.required else '"0..1"'
                lines.append(f'{class_name} {card} --> {req_card} {slot.range} : {slot_name}')
    
    lines.append("")
    lines.append("@enduml")
    
    return '\n'.join(lines)

def main():
    if len(sys.argv) < 2:
        print("Usage: generate_plantuml_modular.py <schema.yaml> [output.puml]")
        print("\nGenerates PlantUML class diagrams from modular LinkML schemas.")
        print("Works around gen-plantuml 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 PlantUML
    print("Generating PlantUML diagram...", file=sys.stderr)
    plantuml = generate_plantuml_from_schemaview(sv)
    
    # Output
    if output_path:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(plantuml)
        print(f"✅ Generated: {output_path}", file=sys.stderr)
        print(f"   Size: {len(plantuml)} bytes", file=sys.stderr)
    else:
        print(plantuml)

if __name__ == '__main__':
    main()