#!/usr/bin/env python3
"""
Test script for the enhanced editing workflow
Tests: Edit Task Plan → Update DAG Visualization → Validate & Deploy
"""

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

from gradio_llm_interface import GradioLlmInterface
import json

def test_task_plan_editor():
    """Test task plan editor functionality"""
    print("Testing Task Plan Editor...")
    print("=" * 40)
    
    # Create test task data
    test_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"]
                }
            }
        ]
    }
    
    interface = GradioLlmInterface()
    
    # Test with existing task plan
    state_with_plan = {'pending_task_plan': test_task_data}
    result = interface.show_task_plan_editor(state_with_plan)
    
    if result and len(result) == 4:
        editor_update, dag_btn_update, validate_btn_update, status_msg = result
        
        print("✓ Editor opened with existing task plan")
        print(f"  Status: {status_msg[:50]}...")
        print(f"  Editor visible: {'✓' if editor_update.get('visible') else '✗'}")
        print(f"  JSON populated: {'✓' if editor_update.get('value') else '✗'}")
        
        return True
    else:
        print("✗ Failed to open task plan editor")
        return False

def test_dag_update_from_editor():
    """Test DAG update from edited JSON"""
    print("\nTesting DAG Update from Editor...")
    print("=" * 40)
    
    # Sample edited JSON
    edited_json = """{
  "tasks": [
    {
      "task": "move_to_position_1",
      "instruction_function": {
        "name": "move_to_position",
        "robot_ids": ["robot_dump_truck_01"],
        "dependencies": [],
        "object_keywords": ["loading_zone"]
      }
    },
    {
      "task": "load_material_1",
      "instruction_function": {
        "name": "load_material",
        "robot_ids": ["robot_dump_truck_01"],
        "dependencies": ["move_to_position_1"],
        "object_keywords": ["soil", "material"]
      }
    }
  ]
}"""
    
    interface = GradioLlmInterface()
    state = {}
    
    try:
        result = interface.update_dag_from_editor(edited_json, state)
        
        if result and len(result) == 6:
            dag_image, validate_btn, editor_vis, dag_btn_vis, status_msg, updated_state = result
            
            print("✓ DAG updated from edited JSON")
            print(f"  Image generated: {'✓' if dag_image else '✗'}")
            print(f"  Validate button shown: {'✓' if validate_btn.get('visible') else '✗'}")
            print(f"  Task plan stored: {'✓' if updated_state.get('pending_task_plan') else '✗'}")
            
            return True
        else:
            print("✗ Failed to update DAG from editor")
            return False
            
    except Exception as e:
        print(f"✗ Error updating DAG: {e}")
        return False

def test_invalid_json_handling():
    """Test handling of invalid JSON"""
    print("\nTesting Invalid JSON Handling...")
    print("=" * 40)
    
    invalid_json = """{
  "tasks": [
    {
      "task": "invalid_task"
      "missing_comma": true
    }
  ]
}"""
    
    interface = GradioLlmInterface()
    state = {}
    
    try:
        result = interface.update_dag_from_editor(invalid_json, state)
        
        if result and len(result) == 6:
            dag_image, validate_btn, editor_vis, dag_btn_vis, status_msg, updated_state = result
            
            if "JSON Parsing Error" in status_msg:
                print("✓ Invalid JSON handled correctly")
                print(f"  Error message displayed: {'✓' if 'JSON Parsing Error' in status_msg else '✗'}")
                print(f"  Editor kept visible: {'✓' if editor_vis.get('visible') else '✗'}")
                return True
            else:
                print("✗ Invalid JSON not handled correctly")
                return False
        else:
            print("✗ Unexpected result from invalid JSON")
            return False
            
    except Exception as e:
        print(f"✗ Unexpected exception: {e}")
        return False

def test_full_workflow():
    """Test the complete editing workflow"""
    print("\nTesting Complete Editing Workflow...")
    print("=" * 40)
    
    interface = GradioLlmInterface()
    
    # Step 1: Open editor
    state = {}
    print("Step 1: Opening task plan editor...")
    editor_result = interface.show_task_plan_editor(state)
    
    if not editor_result:
        print("✗ Failed at step 1")
        return False
    
    # Step 2: Update DAG with valid JSON
    valid_json = """{
  "tasks": [
    {
      "task": "complete_workflow_test",
      "instruction_function": {
        "name": "test_function",
        "robot_ids": ["robot_excavator_01"],
        "dependencies": [],
        "object_keywords": ["test_object"]
      }
    }
  ]
}"""
    
    print("Step 2: Updating DAG from editor...")
    update_result = interface.update_dag_from_editor(valid_json, state)
    
    if not update_result or len(update_result) != 6:
        print("✗ Failed at step 2")
        return False
    
    # Step 3: Validate and deploy
    print("Step 3: Validating and deploying...")
    deploy_result = interface.validate_and_deploy_task_plan(state)
    
    if deploy_result:
        print("✓ Complete workflow test passed")
        return True
    else:
        print("✗ Failed at step 3")
        return False

def main():
    """Run all editing workflow tests"""
    print("🛠️ Enhanced Editing Workflow Tests")
    print("=" * 50)
    
    tests = [
        test_task_plan_editor,
        test_dag_update_from_editor,
        test_invalid_json_handling,
        test_full_workflow
    ]
    
    passed = 0
    total = len(tests)
    
    for test in tests:
        try:
            if test():
                passed += 1
        except Exception as e:
            print(f"✗ Test failed with exception: {e}")
    
    print("\n" + "=" * 50)
    print(f"Editing Workflow Tests passed: {passed}/{total}")
    
    if passed == total:
        print("🎉 All editing workflow tests passed!")
        print("\n🔧 Enhanced Workflow Features:")
        print("  ✓ Manual JSON editing capability")
        print("  ✓ Real-time DAG visualization updates")
        print("  ✓ JSON validation and error handling")
        print("  ✓ Three-step safety workflow:")
        print("    1. 📝 Edit Task Plan")
        print("    2. 🔄 Update DAG Visualization") 
        print("    3. 🔒 Validate & Deploy Task Plan")
        return True
    else:
        print("❌ Some editing workflow tests failed!")
        return False

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