#!/usr/bin/env python3
"""
Test CrewAI integration with the Arzule backend.

This script tests the full flow:
1. Initialize arzule-ingest with API credentials
2. Create a simple CrewAI crew
3. Run it and verify traces are sent to the backend

Usage:
    # Set environment variables first (copy from dashboard Settings > API Keys)
    export ARZULE_API_KEY="arz_live_..."
    export ARZULE_TENANT_ID="..."
    export ARZULE_PROJECT_ID="..."
    
    # Also need an LLM API key for CrewAI
    export ANTHROPIC_API_KEY="sk-ant-..." 
    # or
    export OPENAI_API_KEY="sk-..."
    
    # Run the test
    python test_crewai_integration.py
"""

import os
import sys

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


def load_env():
    """Load environment variables from .env file if it exists."""
    from pathlib import Path
    
    env_path = Path(__file__).parent / ".env"
    if not env_path.exists():
        return
    
    with open(env_path) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, _, value = line.partition("=")
            key = key.strip()
            value = value.strip()
            if key and value and key not in os.environ:
                os.environ[key] = value


def check_requirements():
    """Check that all required environment variables are set."""
    required = ["ARZULE_API_KEY", "ARZULE_TENANT_ID", "ARZULE_PROJECT_ID"]
    missing = [k for k in required if not os.environ.get(k)]
    
    if missing:
        print("Missing required environment variables:")
        for k in missing:
            print(f"  - {k}")
        print()
        print("Set these in your environment or create a .env file.")
        print("You can get these from the Arzule Dashboard: Settings > API Keys")
        return False
    
    # Check for LLM API key
    if not os.environ.get("ANTHROPIC_API_KEY") and not os.environ.get("OPENAI_API_KEY"):
        print("Missing LLM API key.")
        print("Set ANTHROPIC_API_KEY or OPENAI_API_KEY in your environment.")
        return False
    
    return True


def get_llm():
    """Get the LLM based on available API keys."""
    from crewai import LLM
    
    if os.environ.get("ANTHROPIC_API_KEY"):
        print("Using Anthropic Claude")
        return LLM(model="anthropic/claude-3-haiku-20240307")
    elif os.environ.get("OPENAI_API_KEY"):
        print("Using OpenAI GPT-4o-mini")
        return LLM(model="openai/gpt-4o-mini")
    else:
        raise ValueError("No LLM API key found")


def run_test():
    """Run the CrewAI integration test."""
    print("=" * 60)
    print("Arzule CrewAI Integration Test")
    print("=" * 60)
    print()
    
    # Load .env file
    load_env()
    
    # Check requirements
    if not check_requirements():
        sys.exit(1)
    
    # Import CrewAI
    try:
        from crewai import Agent, Crew, Task
    except ImportError:
        print("CrewAI not installed. Run: pip install 'arzule-ingest[crewai]'")
        sys.exit(1)
    
    # Initialize Arzule - this auto-instruments CrewAI
    import arzule_ingest
    
    print()
    print("Initializing Arzule SDK...")
    config = arzule_ingest.init()
    print(f"  Tenant ID: {config['tenant_id']}")
    print(f"  Project ID: {config['project_id']}")
    print(f"  Run ID: {config['run_id']}")
    print(f"  Ingest URL: {config['ingest_url']}")
    print()
    
    # Get LLM
    llm = get_llm()
    
    # Create a simple crew
    print("Creating CrewAI agents and tasks...")
    
    researcher = Agent(
        role="Researcher",
        goal="Research the given topic and provide key insights",
        backstory="You are a skilled researcher who finds accurate information quickly.",
        llm=llm,
        verbose=True,
    )
    
    writer = Agent(
        role="Writer", 
        goal="Write clear and concise summaries",
        backstory="You are an expert writer who creates engaging content.",
        llm=llm,
        verbose=True,
    )
    
    research_task = Task(
        description="Research the benefits of observability in AI agent systems. Provide 3 key points.",
        expected_output="A list of 3 key benefits of observability in AI systems.",
        agent=researcher,
    )
    
    writing_task = Task(
        description="Based on the research findings, write a 2-sentence summary suitable for a tweet.",
        expected_output="A concise 2-sentence summary.",
        agent=writer,
    )
    
    crew = Crew(
        agents=[researcher, writer],
        tasks=[research_task, writing_task],
        verbose=True,
    )
    
    # Run the crew
    print()
    print("-" * 60)
    print("Running crew...")
    print("-" * 60)
    print()
    
    try:
        result = crew.kickoff()
        print()
        print("-" * 60)
        print("Crew completed successfully!")
        print("-" * 60)
        print()
        print("Result:")
        print(result)
    except Exception as e:
        print()
        print("-" * 60)
        print(f"Crew failed with error: {e}")
        print("-" * 60)
    
    # Shutdown to flush events
    print()
    print("Flushing events to Arzule backend...")
    arzule_ingest.shutdown()
    
    print()
    print("=" * 60)
    print("Test complete!")
    print("=" * 60)
    print()
    print("Next steps:")
    print("  1. Go to the Arzule Dashboard (https://arzule.com/dashboard/runs)")
    print("  2. You should see a new run with the captured trace events")
    print("  3. Click on the run to view the trace timeline")
    print()


if __name__ == "__main__":
    run_test()

