#!/usr/bin/env python3
"""
Test Batfish integration and digital twin simulation
"""

import logging
from agent.batfish_client import BatfishClient, BatfishAnalysis
from agent.pipeline_engine import OvergrowthPipeline, NetworkIntent

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def test_batfish_mock_mode():
    """Test Batfish client in mock mode"""
    print("\n=== Test 1: Batfish Mock Mode ===")
    
    client = BatfishClient(use_batfish=True)
    
    # Should be in mock mode without Batfish service
    print(f"Mock mode: {client.mock_mode}")
    
    # Test config analysis
    configs = {
        "router1": """
hostname router1
!
vlan 10
 name Management
!
vlan 20
 name Users
!
interface Vlan10
 ip address 10.0.10.1 255.255.255.0
!
interface GigabitEthernet0/1
 switchport access vlan 10
!
router ospf 1
 network 10.0.0.0 0.255.255.255 area 0
!
""",
        "router2": """
hostname router2
!
vlan 10
 name Management
!
interface Vlan10
 ip address 10.0.10.2 255.255.255.0
!
interface GigabitEthernet0/1
 switchport access vlan 30
!
router ospf 1
 network 10.0.0.0 0.255.255.255 area 0
!
"""
    }
    
    analysis = client.analyze_configs(configs)
    
    print(f"\nAnalysis Results:")
    print(f"  All passed: {analysis.all_passed}")
    print(f"  Undefined references: {len(analysis.undefined_references)}")
    print(f"  Routing loops: {len(analysis.routing_loops)}")
    print(f"  Forwarding errors: {len(analysis.forwarding_errors)}")
    
    if analysis.undefined_references:
        print("\nUndefined References:")
        for ref in analysis.undefined_references:
            print(f"  - {ref['device']}: {ref['type']} {ref['name']}")
    
    recommendations = client.generate_config_recommendations(analysis)
    print(f"\nRecommendations:")
    for rec in recommendations:
        print(f"  - {rec}")
    
    print("\n✓ Mock mode test complete")


def test_acl_validation():
    """Test ACL behavior validation"""
    print("\n=== Test 2: ACL Validation ===")
    
    client = BatfishClient(use_batfish=True)
    
    # Test traffic permissibility
    permitted = client.validate_acl_behavior(
        src="10.0.10.0/24",
        dst="10.0.20.0/24",
        protocol="TCP",
        dst_port=443
    )
    
    print(f"Traffic 10.0.10.0/24 -> 10.0.20.0/24:443/TCP: {'PERMIT' if permitted else 'DENY'}")
    print("✓ ACL validation test complete")


def test_failover_scenario():
    """Test network failover simulation"""
    print("\n=== Test 3: Failover Scenario ===")
    
    client = BatfishClient(use_batfish=True)
    
    # Test if network survives device failure
    survives = client.test_failover_scenario(
        failed_device="core-sw-01",
        src="10.0.10.10",
        dst="10.0.20.20"
    )
    
    print(f"Network survives core-sw-01 failure: {survives}")
    print("✓ Failover test complete")


def test_pipeline_with_batfish():
    """Test pipeline integration with Batfish"""
    print("\n=== Test 4: Pipeline with Batfish ===")
    
    pipeline = OvergrowthPipeline()
    
    # Create network intent
    intent = NetworkIntent(
        description="Enterprise campus network with redundancy",
        business_requirements=["99.99% uptime", "Scalable to 1000 devices"],
        constraints=["Cisco only", "OSPF for routing"],
        budget="$50000",
        timeline="1 month"
    )
    
    # Generate source of truth
    model = pipeline.stage2_generate_sot(intent)
    
    print(f"\nGenerated Network Model:")
    print(f"  Name: {model.name}")
    print(f"  VLANs: {len(model.vlans)}")
    print(f"  Subnets: {len(model.subnets)}")
    print(f"  Devices: {len(model.devices)}")
    
    # Run pre-flight with Batfish
    preflight = pipeline.stage0_preflight(model)
    
    print(f"\nPre-flight Results:")
    print(f"  Schema valid: {preflight['schema_valid']}")
    print(f"  Policy passed: {preflight['policy_passed']}")
    print(f"  Batfish passed: {preflight['batfish_passed']}")
    print(f"  Ready to deploy: {preflight['ready_to_deploy']}")
    
    if 'batfish_analysis' in preflight:
        bf = preflight['batfish_analysis']
        print(f"\nBatfish Analysis:")
        print(f"  Mock mode: {bf.get('mock_mode', False)}")
        print(f"  All passed: {bf.get('all_passed', False)}")
        
        if bf.get('recommendations'):
            print("\nRecommendations:")
            for rec in bf['recommendations'][:3]:
                print(f"  - {rec}")
    
    print("\n✓ Pipeline integration test complete")


def test_config_generation():
    """Test config generation for Batfish"""
    print("\n=== Test 5: Config Generation ===")
    
    from agent.pipeline_engine import NetworkModel, Device, NetworkIntent
    
    # Create a simple network model
    intent = NetworkIntent(
        description="Test network",
        business_requirements=[],
        constraints=[]
    )
    
    model = NetworkModel(
        name="test-network",
        version="1.0.0",
        intent=intent,
        devices=[
            Device(
                name="sw-01",
                role="access",
                model="Catalyst 2960",
                vendor="cisco",
                mgmt_ip="10.0.10.10",
                location="IDF1",
                interfaces=[]
            )
        ],
        vlans=[
            {"id": 10, "name": "Management"},
            {"id": 20, "name": "Users"}
        ],
        subnets=[
            {"network": "10.0.10.0/24", "gateway": "10.0.10.1", "vlan": 10}
        ],
        routing={"protocol": "ospf", "process_id": 1, "networks": ["10.0.0.0/8"]},
        services=["DHCP", "DNS"]
    )
    
    pipeline = OvergrowthPipeline()
    configs = pipeline._generate_configs_for_batfish(model)
    
    print(f"\nGenerated {len(configs)} device configs:")
    for device, config in configs.items():
        print(f"\n{device}:")
        print("-" * 40)
        print(config[:200] + "..." if len(config) > 200 else config)
    
    print("\n✓ Config generation test complete")


def test_real_batfish_if_available():
    """Test with real Batfish service if running"""
    print("\n=== Test 6: Real Batfish (Optional) ===")
    
    import os
    
    # Check if Batfish service is available
    batfish_host = os.getenv("BATFISH_HOST", "localhost")
    
    try:
        import socket
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(2)
        result = sock.connect_ex((batfish_host, 9996))
        sock.close()
        
        if result != 0:
            print("⊘ Batfish service not running - skipping")
            print("  Start Batfish with: docker run -d -p 9996:9996 -p 9997:9997 batfish/batfish")
            return
        
        # Batfish is running - test real analysis
        client = BatfishClient(host=batfish_host, use_batfish=True)
        
        configs = {
            "test-router": """
hostname test-router
interface GigabitEthernet0/0
 ip address 10.0.0.1 255.255.255.0
 no shutdown
!
router ospf 1
 network 10.0.0.0 0.0.0.255 area 0
!
"""
        }
        
        analysis = client.analyze_configs(configs, network_name="real-test")
        
        print(f"✓ Real Batfish analysis complete:")
        print(f"  Mock mode: {client.mock_mode}")
        print(f"  All passed: {analysis.all_passed}")
        
    except ImportError:
        print("⊘ pybatfish not installed - install with: pip install pybatfish")
    except Exception as e:
        print(f"⊘ Batfish test skipped: {e}")


if __name__ == "__main__":
    print("\n" + "="*60)
    print("Batfish Integration & Digital Twin Test Suite")
    print("="*60)
    
    try:
        test_batfish_mock_mode()
        test_acl_validation()
        test_failover_scenario()
        test_config_generation()
        test_pipeline_with_batfish()
        test_real_batfish_if_available()
        
        print("\n" + "="*60)
        print("✓ All Batfish tests completed!")
        print("="*60 + "\n")
        
    except Exception as e:
        print(f"\n✗ Test failed: {e}")
        import traceback
        traceback.print_exc()
        exit(1)