#!/usr/bin/env python
"""
Generate AIDA Workflow Diagram

This script generates a visual diagram of the AIDA agent workflow.
Run from the AIDA directory:
    python scripts/generate_workflow_diagram.py

Requirements:
    pip install pygraphviz  (or graphviz)
"""

import sys
import os

# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.ai.agent.graph import get_aida_graph

def main():
    print("🎨 Generating AIDA Workflow Diagram...")
    
    try:
        # Get the compiled graph
        graph = get_aida_graph()
        
        # Get the underlying graph structure
        drawable = graph.get_graph()
        
        # Get graph name
        graph_name = getattr(graph, 'name', 'AIDA Agent Workflow')
        print(f"📊 Graph: {graph_name}")
        
        # Try different methods to generate the image
        output_path = "agent_workflow.png"
        
        # Method 1: Mermaid PNG (requires playwright or similar)
        try:
            png_data = drawable.draw_mermaid_png()
            with open(output_path, "wb") as f:
                f.write(png_data)
            print(f"✅ Workflow diagram saved to: {output_path}")
            return
        except Exception as e:
            print(f"⚠️ Mermaid PNG failed: {e}")
        
        # Method 2: Try draw_png directly (requires graphviz)
        try:
            png_data = drawable.draw_png()
            with open(output_path, "wb") as f:
                f.write(png_data)
            print(f"✅ Workflow diagram saved to: {output_path}")
            return
        except Exception as e:
            print(f"⚠️ GraphViz PNG failed: {e}")
        
        # Method 3: Generate ASCII art as fallback
        print(f"\n📋 {graph_name} - ASCII Workflow Diagram:")
        print("=" * 70)
        print(drawable.draw_ascii())
        print("=" * 70)
        
        # Method 4: Generate Mermaid text that can be pasted into mermaid.live
        mermaid_code = drawable.draw_mermaid()
        
        # Add title to mermaid code
        mermaid_with_title = mermaid_code.replace(
            "%%{init:",
            f"---\ntitle: {graph_name}\n---\n%%{{init:"
        )
        
        print("\n📝 Mermaid Diagram Code (paste into https://mermaid.live):")
        print("-" * 70)
        print(mermaid_with_title)
        print("-" * 70)
        
        # Save mermaid code to file
        with open("agent_workflow.mmd", "w") as f:
            f.write(mermaid_with_title)
        print(f"✅ Mermaid code saved to: agent_workflow.mmd")
        
    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()