#!/usr/bin/env python3
"""
Test script for the stage function implementation.

This script tests the stage detection and counting functionality without requiring
the full NodeGraphQt dependency.
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Test the core pipeline functions directly
def get_stage_count(node_graph):
    """Mock version of get_stage_count for testing."""
    if not node_graph:
        return 0
    
    all_nodes = node_graph.all_nodes()
    model_nodes = [node for node in all_nodes if 'model' in node.node_type]
    
    return len(model_nodes)

def get_pipeline_summary(node_graph):
    """Mock version of get_pipeline_summary for testing."""
    if not node_graph:
        return {'stage_count': 0, 'valid': False, 'error': 'No pipeline graph'}
    
    all_nodes = node_graph.all_nodes()
    model_nodes = [node for node in all_nodes if 'model' in node.node_type]
    input_nodes = [node for node in all_nodes if 'input' in node.node_type]
    output_nodes = [node for node in all_nodes if 'output' in node.node_type]
    
    # Basic validation
    valid = len(input_nodes) > 0 and len(output_nodes) > 0 and len(model_nodes) > 0
    error = None
    
    if not input_nodes:
        error = "No input nodes found"
    elif not output_nodes:
        error = "No output nodes found"
    elif not model_nodes:
        error = "No model nodes found"
    
    return {
        'stage_count': len(model_nodes),
        'valid': valid,
        'error': error,
        'total_nodes': len(all_nodes),
        'input_nodes': len(input_nodes),
        'output_nodes': len(output_nodes),
        'model_nodes': len(model_nodes),
        'preprocess_nodes': len([n for n in all_nodes if 'preprocess' in n.node_type]),
        'postprocess_nodes': len([n for n in all_nodes if 'postprocess' in n.node_type]),
        'stages': []
    }


class MockPort:
    """Mock port for testing without NodeGraphQt."""
    def __init__(self, node, port_type):
        self.node_ref = node
        self.port_type = port_type
        self.connections = []
    
    def node(self):
        return self.node_ref
    
    def connected_inputs(self):
        return [conn for conn in self.connections if conn.port_type == 'input']
    
    def connected_outputs(self):
        return [conn for conn in self.connections if conn.port_type == 'output']


class MockNode:
    """Mock node for testing without NodeGraphQt."""
    def __init__(self, node_type):
        self.node_type = node_type
        self.input_ports = []
        self.output_ports = []
        self.node_name = f"{node_type}_node"
        self.node_id = f"{node_type}_{id(self)}"
    
    def inputs(self):
        return self.input_ports
    
    def outputs(self):
        return self.output_ports
    
    def add_input(self, name):
        port = MockPort(self, 'input')
        self.input_ports.append(port)
        return port
    
    def add_output(self, name):
        port = MockPort(self, 'output')
        self.output_ports.append(port)
        return port
    
    def name(self):
        return self.node_name


class MockNodeGraph:
    """Mock node graph for testing without NodeGraphQt."""
    def __init__(self):
        self.nodes = []
    
    def all_nodes(self):
        return self.nodes
    
    def add_node(self, node):
        self.nodes.append(node)
    
    def connect_nodes(self, output_node, input_node):
        """Connect output of first node to input of second node."""
        output_port = output_node.add_output('output')
        input_port = input_node.add_input('input')
        
        # Create bidirectional connection
        output_port.connections.append(input_port)
        input_port.connections.append(output_port)


def create_mock_pipeline():
    """Create a mock pipeline for testing."""
    graph = MockNodeGraph()
    
    # Create nodes
    input_node = MockNode('input')
    preprocess_node = MockNode('preprocess')
    model_node1 = MockNode('model')
    postprocess_node1 = MockNode('postprocess')
    model_node2 = MockNode('model')
    postprocess_node2 = MockNode('postprocess')
    output_node = MockNode('output')
    
    # Add nodes to graph
    for node in [input_node, preprocess_node, model_node1, postprocess_node1, 
                model_node2, postprocess_node2, output_node]:
        graph.add_node(node)
    
    # Connect nodes: input -> preprocess -> model1 -> postprocess1 -> model2 -> postprocess2 -> output
    graph.connect_nodes(input_node, preprocess_node)
    graph.connect_nodes(preprocess_node, model_node1)
    graph.connect_nodes(model_node1, postprocess_node1)
    graph.connect_nodes(postprocess_node1, model_node2)
    graph.connect_nodes(model_node2, postprocess_node2)
    graph.connect_nodes(postprocess_node2, output_node)
    
    return graph


def test_stage_count():
    """Test the stage counting functionality."""
    print("Testing Stage Count Function...")
    
    # Create mock pipeline
    graph = create_mock_pipeline()
    
    # Count stages - should be 2 (2 model nodes)
    stage_count = get_stage_count(graph)
    print(f"Stage count: {stage_count}")
    
    # Expected: 2 stages (2 model nodes)
    assert stage_count == 2, f"Expected 2 stages, got {stage_count}"
    print("✓ Stage count test passed")


def test_empty_pipeline():
    """Test with empty pipeline."""
    print("\nTesting Empty Pipeline...")
    
    empty_graph = MockNodeGraph()
    stage_count = get_stage_count(empty_graph)
    print(f"Empty pipeline stage count: {stage_count}")
    
    assert stage_count == 0, f"Expected 0 stages, got {stage_count}"
    print("✓ Empty pipeline test passed")


def test_single_stage():
    """Test with single stage pipeline."""
    print("\nTesting Single Stage Pipeline...")
    
    graph = MockNodeGraph()
    
    # Create simple pipeline: input -> model -> output
    input_node = MockNode('input')
    model_node = MockNode('model')
    output_node = MockNode('output')
    
    graph.add_node(input_node)
    graph.add_node(model_node)
    graph.add_node(output_node)
    
    graph.connect_nodes(input_node, model_node)
    graph.connect_nodes(model_node, output_node)
    
    stage_count = get_stage_count(graph)
    print(f"Single stage pipeline count: {stage_count}")
    
    assert stage_count == 1, f"Expected 1 stage, got {stage_count}"
    print("✓ Single stage test passed")


def test_pipeline_summary():
    """Test the pipeline summary function."""
    print("\nTesting Pipeline Summary...")
    
    graph = create_mock_pipeline()
    
    # Get summary
    summary = get_pipeline_summary(graph)
    
    print(f"Pipeline summary: {summary}")
    
    # Check basic structure
    assert 'stage_count' in summary, "Missing stage_count in summary"
    assert 'valid' in summary, "Missing valid in summary"
    assert 'total_nodes' in summary, "Missing total_nodes in summary"
    
    # Check values
    assert summary['stage_count'] == 2, f"Expected 2 stages, got {summary['stage_count']}"
    assert summary['total_nodes'] == 7, f"Expected 7 nodes, got {summary['total_nodes']}"
    
    print("✓ Pipeline summary test passed")


def main():
    """Run all tests."""
    print("Running Stage Function Tests...")
    print("=" * 50)
    
    try:
        test_stage_count()
        test_empty_pipeline()
        test_single_stage()
        test_pipeline_summary()
        
        print("\n" + "=" * 50)
        print("All tests passed! ✓")
        
    except Exception as e:
        print(f"\n❌ Test failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == '__main__':
    main()