#!/usr/bin/env python3
"""
Test script for DAG visualization functionality
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from dag_visualizer import DAGVisualizer
import json

def test_single_task():
    """Test with a single task (no dependencies)"""
    print("Testing single task visualization...")
    
    task_data = {
        "tasks": [
            {
                "task": "target_area_for_specific_robots_1",
                "instruction_function": {
                    "name": "target_area_for_specific_robots",
                    "robot_ids": ["robot_dump_truck_01"],
                    "dependencies": [],
                    "object_keywords": ["puddle1"]
                }
            }
        ]
    }
    
    visualizer = DAGVisualizer()
    image_path = visualizer.create_dag_visualization(task_data)
    
    if image_path and os.path.exists(image_path):
        print(f"✓ Single task visualization created: {image_path}")
        return True
    else:
        print("✗ Failed to create single task visualization")
        return False

def test_multiple_tasks_with_dependencies():
    """Test with multiple tasks with dependencies"""
    print("Testing multiple tasks with dependencies...")
    
    task_data = {
        "tasks": [
            {
                "task": "target_area_for_specific_robots_1",
                "instruction_function": {
                    "name": "target_area_for_specific_robots",
                    "robot_ids": ["robot_dump_truck_01"],
                    "dependencies": [],
                    "object_keywords": ["puddle1"]
                }
            },
            {
                "task": "avoid_areas_for_all_robots_1",
                "instruction_function": {
                    "name": "avoid_areas_for_all_robots",
                    "robot_ids": ["robot_dump_truck_01", "robot_excavator_01"],
                    "dependencies": ["target_area_for_specific_robots_1"],
                    "object_keywords": ["obstacle1", "obstacle2"]
                }
            },
            {
                "task": "move_to_position_1",
                "instruction_function": {
                    "name": "move_to_position",
                    "robot_ids": ["robot_excavator_01"],
                    "dependencies": ["avoid_areas_for_all_robots_1"],
                    "object_keywords": ["soil_pile"]
                }
            }
        ]
    }
    
    visualizer = DAGVisualizer()
    image_path = visualizer.create_dag_visualization(task_data)
    
    if image_path and os.path.exists(image_path):
        print(f"✓ Multiple tasks visualization created: {image_path}")
        return True
    else:
        print("✗ Failed to create multiple tasks visualization")
        return False

def test_simplified_visualization():
    """Test simplified visualization"""
    print("Testing simplified visualization...")
    
    task_data = {
        "tasks": [
            {
                "task": "excavate_soil_from_pile",
                "instruction_function": {
                    "name": "excavate_soil_from_pile",
                    "robot_ids": ["robot_excavator_01"],
                    "dependencies": [],
                    "object_keywords": ["soil_pile"]
                }
            },
            {
                "task": "transport_soil_to_pit",
                "instruction_function": {
                    "name": "transport_soil_to_pit",
                    "robot_ids": ["robot_dump_truck_01"],
                    "dependencies": ["excavate_soil_from_pile"],
                    "object_keywords": ["pit"]
                }
            }
        ]
    }
    
    visualizer = DAGVisualizer()
    image_path = visualizer.create_simplified_dag_visualization(task_data)
    
    if image_path and os.path.exists(image_path):
        print(f"✓ Simplified visualization created: {image_path}")
        return True
    else:
        print("✗ Failed to create simplified visualization")
        return False

def test_invalid_data():
    """Test with invalid data"""
    print("Testing invalid data handling...")
    
    # Test with empty data
    visualizer = DAGVisualizer()
    image_path = visualizer.create_dag_visualization({})
    
    if image_path is None:
        print("✓ Invalid data handled correctly (returned None)")
        return True
    else:
        print("✗ Invalid data not handled correctly")
        return False

def main():
    """Run all tests"""
    print("Starting DAG Visualization Tests...")
    print("=" * 50)
    
    tests = [
        test_single_task,
        test_multiple_tasks_with_dependencies,
        test_simplified_visualization,
        test_invalid_data
    ]
    
    passed = 0
    total = len(tests)
    
    for test in tests:
        try:
            if test():
                passed += 1
            print()
        except Exception as e:
            print(f"✗ Test failed with exception: {e}")
            print()
    
    print("=" * 50)
    print(f"Tests passed: {passed}/{total}")
    
    if passed == total:
        print("🎉 All tests passed!")
        return True
    else:
        print("❌ Some tests failed!")
        return False

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)