#!/usr/bin/env python3
"""
Test the live pipeline with an abandoned handoff scenario.

This sends events that should:
1. Be ingested by arzule-ingest
2. Trigger arzule-forensics on run.end
3. Detect the abandoned handoff
4. Trigger arzule-advisor to generate suggestions
"""

import os
import sys
import uuid

# Add src to path for local testing
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))

from arzule_ingest.sinks.http_batch import HttpBatchSink
from arzule_ingest.run import ArzuleRun
from arzule_ingest.ids import new_span_id

# Configuration
INGEST_URL = "https://uuczh0e8g5.execute-api.us-east-1.amazonaws.com/ingest"
API_KEY = os.environ.get("ARZULE_API_KEY", "test-key-123")
TENANT_ID = os.environ.get("ARZULE_TENANT_ID", "test-tenant")
PROJECT_ID = os.environ.get("ARZULE_PROJECT_ID", "test-project")


def test_abandoned_handoff():
    """Send a trace with an abandoned handoff to test the full pipeline."""
    
    print(f"Testing live pipeline at: {INGEST_URL}")
    print(f"Tenant: {TENANT_ID}, Project: {PROJECT_ID}")
    print("-" * 60)
    
    # Create HTTP sink
    sink = HttpBatchSink(
        endpoint_url=INGEST_URL,
        api_key=API_KEY,
        require_tls=True,
    )
    
    # Create a run
    run = ArzuleRun(
        tenant_id=TENANT_ID,
        project_id=PROJECT_ID,
        sink=sink,
    )
    
    with run:
        print(f"Run ID: {run.run_id}")
        print(f"Trace ID: {run.trace_id}")
        
        # === Scenario: Abandoned Handoff ===
        # Agent A proposes a handoff to Agent B, but B never acknowledges it
        
        handoff_key = str(uuid.uuid4())
        
        # Agent A starts
        agent_a_span = new_span_id()
        run.emit(run._make_event(
            event_type="agent.execution.start",
            span_id=agent_a_span,
            parent_span_id=run._root_span_id,
            summary="Agent Researcher started",
            agent={"id": "crewai:role:Researcher", "role": "Researcher"},
        ))
        print("  [1] agent.execution.start - Researcher")
        
        # Agent A does some work (LLM call)
        llm_span = new_span_id()
        run.emit(run._make_event(
            event_type="llm.start",
            span_id=llm_span,
            parent_span_id=agent_a_span,
            summary="LLM call started",
            agent={"id": "crewai:role:Researcher", "role": "Researcher"},
            attrs_compact={"model": "gpt-4", "provider": "openai"},
        ))
        print("  [2] llm.start")
        
        run.emit(run._make_event(
            event_type="llm.end",
            span_id=llm_span,
            parent_span_id=agent_a_span,
            summary="LLM call completed",
            agent={"id": "crewai:role:Researcher", "role": "Researcher"},
            attrs_compact={"model": "gpt-4", "tokens_total": 150},
        ))
        print("  [3] llm.end")
        
        # Agent A proposes handoff to Writer (BUT WRITER NEVER ACKS!)
        handoff_span = new_span_id()
        run.emit(run._make_event(
            event_type="handoff.proposed",
            span_id=handoff_span,
            parent_span_id=agent_a_span,
            summary="handoff proposed to Writer",
            agent={"id": "crewai:role:Researcher", "role": "Researcher"},
            attrs_compact={
                "handoff_key": handoff_key,
                "from_agent_role": "Researcher",
                "to_coworker": "Writer",
                "tool_name": "delegate_work_to_coworker",
            },
            payload={
                "tool_input": {
                    "task": "Write a summary of the research findings",
                    "coworker": "Writer",
                    "context": "Found 3 key insights about AI safety",
                }
            },
        ))
        print(f"  [4] handoff.proposed (key={handoff_key[:8]}...)")
        
        # Agent A completes WITHOUT the handoff being acknowledged
        # This is the failure condition - abandoned handoff
        run.emit(run._make_event(
            event_type="agent.execution.complete",
            span_id=agent_a_span,
            parent_span_id=run._root_span_id,
            status="ok",  # Agent thinks it succeeded
            summary="Agent Researcher completed",
            agent={"id": "crewai:role:Researcher", "role": "Researcher"},
            attrs_compact={
                "handoff_abandoned": handoff_key,
                "expected_receiver": "Writer",
            },
        ))
        print("  [5] agent.execution.complete - Researcher (handoff ABANDONED!)")
        
    # run.__exit__ sends run.end and flushes
    print("-" * 60)
    print("Run completed. Events sent to ingest endpoint.")
    print()
    print("What should happen next (async):")
    print("  1. arzule-ingest stores events in trace_events_index")
    print("  2. run.end triggers arzule-forensics")
    print("  3. forensics detects abandoned handoff -> writes to findings table")
    print("  4. forensics invokes arzule-advisor")
    print("  5. advisor calls Bedrock Claude -> writes to suggestions table")
    print()
    print("Check the AWS tables to verify:")
    print(f"  - findings table: should have ABANDONED_HANDOFF for run {run.run_id}")
    print(f"  - suggestions table: should have LLM fix suggestions for run {run.run_id}")
    
    sink.close()
    
    return run.run_id


if __name__ == "__main__":
    run_id = test_abandoned_handoff()
    print()
    print(f"Test complete. Run ID: {run_id}")




