"""
Test all 11 forensic detectors end-to-end.

Each function sends a trace designed to trigger a specific detector.
"""
import os
import uuid
import time
from dotenv import load_dotenv

load_dotenv()

import arzule_ingest as arzule


def send_trace(name: str, emit_events):
    """Helper to send a trace and return the run_id."""
    config = arzule.init()
    run_id = config["run_id"]
    run = arzule.current_run()
    
    print(f"\n[{name}] Run ID: {run_id}")
    
    emit_events(run)
    
    arzule.shutdown()
    
    # Reset for next test
    arzule._initialized = False
    arzule._global_sink = None
    arzule._global_run = None
    
    return run_id


def test_1_abandoned_handoff():
    """AbandonedHandoffDetector: Handoff proposed but never acknowledged."""
    def emit(run):
        handoff_key = str(uuid.uuid4())
        run.emit(run._make_event(
            event_type='handoff.proposed',
            agent={'id': 'crewai:role:Planner', 'role': 'Planner'},
            status='ok',
            summary='Handoff to Writer',
            attrs_compact={
                'handoff_key': handoff_key,
                'to_agent': {'id': 'crewai:role:Writer', 'role': 'Writer'},
            },
        ))
        # No handoff.ack or handoff.complete
    
    return send_trace("AbandonedHandoff", emit)


def test_2_missing_ownership():
    """MissingOwnershipDetector: End event without matching start."""
    def emit(run):
        # Task complete without task start
        run.emit(run._make_event(
            event_type='task.complete',
            agent={'id': 'crewai:role:Worker', 'role': 'Worker'},
            status='ok',
            summary='Task completed',
            task_id='orphan-task-123',
            attrs_compact={'orphan': True},
        ))
    
    return send_trace("MissingOwnership", emit)


def test_3_circular_delegation():
    """CircularDelegationDetector: A -> B -> A delegation cycle."""
    def emit(run):
        # A proposes to B - need from_agent_role and to_coworker for edge building
        key1 = str(uuid.uuid4())
        run.emit(run._make_event(
            event_type='handoff.proposed',
            agent={'id': 'crewai:role:AgentA', 'role': 'AgentA'},
            status='ok',
            attrs_compact={
                'handoff_key': key1,
                'from_agent_role': 'AgentA',  # Explicit from
                'to_coworker': 'AgentB',  # Used by graph builder
                'to_agent': {'id': 'crewai:role:AgentB', 'role': 'AgentB'},
            },
        ))
        run.emit(run._make_event(
            event_type='handoff.ack',
            agent={'id': 'crewai:role:AgentB', 'role': 'AgentB'},
            status='ok',
            attrs_compact={
                'handoff_key': key1,
                'to_agent_role': 'AgentB',
            },
        ))
        
        # B proposes back to A (cycle!)
        key2 = str(uuid.uuid4())
        run.emit(run._make_event(
            event_type='handoff.proposed',
            agent={'id': 'crewai:role:AgentB', 'role': 'AgentB'},
            status='ok',
            attrs_compact={
                'handoff_key': key2,
                'from_agent_role': 'AgentB',
                'to_coworker': 'AgentA',
                'to_agent': {'id': 'crewai:role:AgentA', 'role': 'AgentA'},
            },
        ))
        run.emit(run._make_event(
            event_type='handoff.ack',
            agent={'id': 'crewai:role:AgentA', 'role': 'AgentA'},
            status='ok',
            attrs_compact={
                'handoff_key': key2,
                'to_agent_role': 'AgentA',
            },
        ))
    
    return send_trace("CircularDelegation", emit)


def test_4_handoff_mismatch():
    """HandoffMismatchDetector: Payload sent differs from received - explicit mismatch flag."""
    def emit(run):
        handoff_key = str(uuid.uuid4())
        # Proposed with tool_input in attrs_compact (this gets indexed)
        run.emit(run._make_event(
            event_type='handoff.proposed',
            agent={'id': 'crewai:role:Sender', 'role': 'Sender'},
            status='ok',
            attrs_compact={
                'handoff_key': handoff_key,
                'from_agent_role': 'Sender',
                'to_coworker': 'Receiver',
                'to_agent': {'id': 'crewai:role:Receiver', 'role': 'Receiver'},
                'tool_input': {'task': 'write article', 'topic': 'AI'},  # Proposed input
            },
        ))
        # Ack with explicit payload_mismatch flag
        run.emit(run._make_event(
            event_type='handoff.ack',
            agent={'id': 'crewai:role:Receiver', 'role': 'Receiver'},
            status='ok',
            attrs_compact={
                'handoff_key': handoff_key,
                'to_agent_role': 'Receiver',
                'payload_mismatch': True,  # Explicit mismatch flag
            },
        ))
    
    return send_trace("HandoffMismatch", emit)


def test_5_cascading_error():
    """CascadingErrorDetector: Root error causes downstream failures."""
    def emit(run):
        # Root error with is_root_cause and error_id
        root_error_id = str(uuid.uuid4())
        run.emit(run._make_event(
            event_type='tool.call.error',
            agent={'id': 'crewai:role:DataFetcher', 'role': 'DataFetcher'},
            status='error',
            summary='API call failed',
            attrs_compact={
                'tool_name': 'fetch_data',
                'error_type': 'ConnectionError',
                'error_message': 'Connection refused',
                'error_id': root_error_id,  # ID for linking
                'is_root_cause': True,  # Explicit marker
            },
        ))
        
        # Downstream failure 1 - linked to root
        run.emit(run._make_event(
            event_type='agent.execution.error',
            agent={'id': 'crewai:role:Analyzer', 'role': 'Analyzer'},
            status='error',
            summary='Analysis failed - no data',
            attrs_compact={
                'caused_by_error_id': root_error_id,  # Links to root
            },
        ))
        # Downstream failure 2 - linked to root
        run.emit(run._make_event(
            event_type='agent.execution.error',
            agent={'id': 'crewai:role:Reporter', 'role': 'Reporter'},
            status='error',
            summary='Report failed - no analysis',
            attrs_compact={
                'caused_by_error_id': root_error_id,  # Also links to root
            },
        ))
    
    return send_trace("CascadingError", emit)


def test_6_failed_handoff():
    """FailedHandoffDetector: Handoff acknowledged but completed with error."""
    def emit(run):
        handoff_key = str(uuid.uuid4())
        run.emit(run._make_event(
            event_type='handoff.proposed',
            agent={'id': 'crewai:role:Manager', 'role': 'Manager'},
            status='ok',
            attrs_compact={
                'handoff_key': handoff_key,
                'to_agent': {'id': 'crewai:role:Worker', 'role': 'Worker'},
            },
        ))
        run.emit(run._make_event(
            event_type='handoff.ack',
            agent={'id': 'crewai:role:Worker', 'role': 'Worker'},
            status='ok',
            attrs_compact={'handoff_key': handoff_key},
        ))
        run.emit(run._make_event(
            event_type='handoff.complete',
            agent={'id': 'crewai:role:Worker', 'role': 'Worker'},
            status='error',  # Failed!
            summary='Task failed',
            attrs_compact={'handoff_key': handoff_key, 'error': 'Timeout'},
        ))
    
    return send_trace("FailedHandoff", emit)


def test_7_tool_error():
    """ToolErrorDetector: Tool execution failures."""
    def emit(run):
        run.emit(run._make_event(
            event_type='tool.call.start',
            agent={'id': 'crewai:role:Researcher', 'role': 'Researcher'},
            status='ok',
            attrs_compact={'tool_name': 'web_search'},
        ))
        run.emit(run._make_event(
            event_type='tool.call.error',
            agent={'id': 'crewai:role:Researcher', 'role': 'Researcher'},
            status='error',
            summary='Tool failed',
            attrs_compact={
                'tool_name': 'web_search',
                'error': 'Rate limit exceeded',
                'retry_count': 3,
            },
        ))
    
    return send_trace("ToolError", emit)


def test_8_context_drift():
    """ContextDriftDetector: Agent output contradicts input assumptions via explicit marker."""
    def emit(run):
        # Need a complete handoff for context drift detection
        handoff_key = str(uuid.uuid4())
        run.emit(run._make_event(
            event_type='handoff.proposed',
            agent={'id': 'crewai:role:Manager', 'role': 'Manager'},
            status='ok',
            attrs_compact={
                'handoff_key': handoff_key,
                'from_agent_role': 'Manager',
                'to_coworker': 'Writer',
                'to_agent': {'id': 'crewai:role:Writer', 'role': 'Writer'},
            },
            payload={
                'tool_input': {'task': 'Write about Python', 'context': 'Focus on Python 3.11 features'},
            },
        ))
        run.emit(run._make_event(
            event_type='handoff.ack',
            agent={'id': 'crewai:role:Writer', 'role': 'Writer'},
            status='ok',
            attrs_compact={
                'handoff_key': handoff_key,
                'to_agent_role': 'Writer',
            },
        ))
        # Complete with drift marker
        run.emit(run._make_event(
            event_type='handoff.complete',
            agent={'id': 'crewai:role:Writer', 'role': 'Writer'},
            status='ok',
            attrs_compact={
                'handoff_key': handoff_key,
                'context_drift': True,  # Explicit drift marker
            },
            payload={
                'result': 'Article about Java instead of Python',  # Wrong topic
            },
        ))
    
    return send_trace("ContextDrift", emit)


def test_9_stale_assumption():
    """StaleAssumptionDetector: Agent uses outdated information via explicit marker or cache age."""
    def emit(run):
        # Use explicit stale_data marker
        run.emit(run._make_event(
            event_type='agent.execution.start',
            agent={'id': 'crewai:role:Analyzer', 'role': 'Analyzer'},
            status='ok',
        ))
        run.emit(run._make_event(
            event_type='tool.call.end',
            agent={'id': 'crewai:role:Analyzer', 'role': 'Analyzer'},
            status='ok',
            attrs_compact={
                'tool_name': 'fetch_report',
                'stale_data': 'Report is from 2 days ago',  # Explicit stale marker
            },
        ))
        # Also test cache age threshold
        run.emit(run._make_event(
            event_type='tool.call.end',
            agent={'id': 'crewai:role:Analyzer', 'role': 'Analyzer'},
            status='ok',
            attrs_compact={
                'tool_name': 'get_cached_data',
                'cache_age_seconds': 600,  # 10 minutes > 5 min threshold
            },
        ))
        run.emit(run._make_event(
            event_type='agent.execution.complete',
            agent={'id': 'crewai:role:Analyzer', 'role': 'Analyzer'},
            status='ok',
        ))
    
    return send_trace("StaleAssumption", emit)


def test_10_broken_verification():
    """BrokenVerificationGateDetector: Skipped or failed verification."""
    def emit(run):
        run.emit(run._make_event(
            event_type='agent.execution.start',
            agent={'id': 'crewai:role:QA', 'role': 'QA'},
            status='ok',
            attrs_compact={'verification_required': True},
        ))
        run.emit(run._make_event(
            event_type='agent.execution.complete',
            agent={'id': 'crewai:role:QA', 'role': 'QA'},
            status='ok',
            attrs_compact={
                'verification_performed': False,  # Skipped!
                'verification_skipped': True,
                'reason': 'Timeout pressure',
            },
        ))
    
    return send_trace("BrokenVerification", emit)


def test_11_state_desync():
    """StateDesyncDetector: State mismatch during handoff."""
    def emit(run):
        handoff_key = str(uuid.uuid4())
        run.emit(run._make_event(
            event_type='handoff.proposed',
            agent={'id': 'crewai:role:Producer', 'role': 'Producer'},
            status='ok',
            attrs_compact={
                'handoff_key': handoff_key,
                'to_agent': {'id': 'crewai:role:Consumer', 'role': 'Consumer'},
                'state': {'items': ['a', 'b', 'c'], 'count': 3},
            },
        ))
        run.emit(run._make_event(
            event_type='handoff.ack',
            agent={'id': 'crewai:role:Consumer', 'role': 'Consumer'},
            status='ok',
            attrs_compact={
                'handoff_key': handoff_key,
                'received_state': {'items': ['a', 'b'], 'count': 2},  # Desync!
                'state_desync': True,
            },
        ))
    
    return send_trace("StateDesync", emit)


def main():
    print("=" * 60)
    print("TESTING ALL 11 DETECTORS")
    print("=" * 60)
    
    run_ids = []
    
    tests = [
        ("1. AbandonedHandoff", test_1_abandoned_handoff),
        ("2. MissingOwnership", test_2_missing_ownership),
        ("3. CircularDelegation", test_3_circular_delegation),
        ("4. HandoffMismatch", test_4_handoff_mismatch),
        ("5. CascadingError", test_5_cascading_error),
        ("6. FailedHandoff", test_6_failed_handoff),
        ("7. ToolError", test_7_tool_error),
        ("8. ContextDrift", test_8_context_drift),
        ("9. StaleAssumption", test_9_stale_assumption),
        ("10. BrokenVerification", test_10_broken_verification),
        ("11. StateDesync", test_11_state_desync),
    ]
    
    for name, test_fn in tests:
        try:
            run_id = test_fn()
            run_ids.append((name, run_id))
            print(f"  -> Sent successfully")
        except Exception as e:
            print(f"  -> ERROR: {e}")
            run_ids.append((name, f"ERROR: {e}"))
    
    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)
    for name, run_id in run_ids:
        print(f"{name}: {run_id}")
    
    print("\nWait 10 seconds for forensics to process...")
    time.sleep(10)
    
    print("\nRun IDs for verification:")
    for name, run_id in run_ids:
        if not run_id.startswith("ERROR"):
            print(f"  {run_id}")


if __name__ == "__main__":
    main()

