#!/usr/bin/env python3
"""Test script for LangChain integration with Arzule.

This script tests the LangChain callback handler by running a simple
chain and agent scenario, verifying that trace events are properly captured.

Usage:
    # Create venv and install dependencies
    python -m venv .venv-langchain-test
    source .venv-langchain-test/bin/activate
    pip install -e ".[langchain]"
    pip install langchain-openai

    # Set your OpenAI API key
    export OPENAI_API_KEY=your-key-here

    # Run the test
    python test_langchain_integration.py
"""

import json
import os
import sys
from pathlib import Path

# Add src to path for development
sys.path.insert(0, str(Path(__file__).parent / "src"))


def test_callback_handler_basic():
    """Test basic callback handler functionality without LLM calls."""
    print("\n=== Test 1: Basic Callback Handler (no LLM) ===\n")

    from arzule_ingest.langchain import instrument_langchain
    from arzule_ingest.langchain.callback_handler import ArzuleLangChainHandler
    from arzule_ingest import ArzuleRun
    from arzule_ingest.sinks.file_jsonl import JsonlFileSink

    # Create handler
    handler = instrument_langchain()
    assert isinstance(handler, ArzuleLangChainHandler)
    print("[OK] Handler created successfully")

    # Check feature flags
    assert handler.enable_llm is True
    assert handler.enable_chain is True
    assert handler.enable_tool is True
    assert handler.enable_agent is True
    print("[OK] Feature flags set correctly")

    print("[OK] Basic callback handler test passed\n")


def test_event_emission_mock():
    """Test event emission with mock objects (no actual LLM calls)."""
    print("\n=== Test 2: Event Emission with Mocks ===\n")

    from uuid import uuid4
    from arzule_ingest.langchain import instrument_langchain
    from arzule_ingest import ArzuleRun
    from arzule_ingest.sinks.file_jsonl import JsonlFileSink

    # Create output file
    output_file = Path("out/langchain_test_mock.jsonl")
    output_file.parent.mkdir(exist_ok=True)

    # Clean up previous run
    if output_file.exists():
        output_file.unlink()

    sink = JsonlFileSink(str(output_file))
    handler = instrument_langchain()

    with ArzuleRun(tenant_id="test-tenant", project_id="test-project", sink=sink) as run:
        # Simulate LLM start
        handler.on_llm_start(
            serialized={"id": ["langchain", "llms", "openai", "ChatOpenAI"], "kwargs": {"model_name": "gpt-4"}},
            prompts=["What is the capital of France?"],
            run_id=uuid4(),
            tags=["test"],
            metadata={"test_run": True},
        )

        # Simulate LLM end (mock response)
        class MockGeneration:
            text = "The capital of France is Paris."

        class MockLLMResult:
            generations = [[MockGeneration()]]
            llm_output = {"token_usage": {"total_tokens": 50, "prompt_tokens": 20, "completion_tokens": 30}}

        handler.on_llm_end(
            response=MockLLMResult(),
            run_id=uuid4(),
        )

        # Simulate chain start
        chain_run_id = uuid4()
        handler.on_chain_start(
            serialized={"id": ["langchain", "chains", "LLMChain"], "name": "TestChain"},
            inputs={"input": "test input"},
            run_id=chain_run_id,
            tags=["chain-test"],
            metadata={"chain_type": "llm"},
        )

        # Simulate chain end
        handler.on_chain_end(
            outputs={"output": "test output"},
            run_id=chain_run_id,
        )

        # Simulate agent action
        class MockAgentAction:
            tool = "search"
            tool_input = {"query": "weather in Paris"}
            log = "Thinking: I should search for the weather"

        handler.on_agent_action(
            action=MockAgentAction(),
            run_id=uuid4(),
            tags=["agent:ResearchAgent"],
            metadata={"node_name": "researcher"},
        )

        # Simulate agent finish
        class MockAgentFinish:
            return_values = {"output": "The weather in Paris is sunny."}
            log = "Final answer provided"

        handler.on_agent_finish(
            finish=MockAgentFinish(),
            run_id=uuid4(),
            tags=["agent:ResearchAgent"],
            metadata={"node_name": "researcher"},
        )

    # Verify output
    if output_file.exists():
        events = []
        with open(output_file) as f:
            for line in f:
                if line.strip():
                    events.append(json.loads(line))

        print(f"[OK] Generated {len(events)} events")

        # Check event types
        event_types = [e.get("event_type") for e in events]
        print(f"    Event types: {event_types}")

        # Verify key events are present
        assert "run.start" in event_types, "Missing run.start"
        assert "llm.call.start" in event_types, "Missing llm.call.start"
        assert "llm.call.end" in event_types, "Missing llm.call.end"
        assert "chain.start" in event_types, "Missing chain.start"
        assert "chain.end" in event_types, "Missing chain.end"
        assert "agent.action" in event_types, "Missing agent.action"
        assert "agent.finish" in event_types, "Missing agent.finish"
        assert "run.end" in event_types, "Missing run.end"

        # Verify agent name extraction
        agent_events = [e for e in events if e.get("event_type") in ("agent.action", "agent.finish")]
        for evt in agent_events:
            agent_info = evt.get("agent", {})
            print(f"    Agent event: {evt['event_type']}, agent={agent_info}")
            assert agent_info.get("role") == "researcher", f"Expected 'researcher', got {agent_info.get('role')}"

        print("[OK] Event emission test passed\n")
    else:
        print("[FAIL] No output file created")
        sys.exit(1)


def test_with_real_langchain():
    """Test with actual LangChain components (requires langchain installed)."""
    print("\n=== Test 3: Real LangChain Integration ===\n")

    try:
        from langchain_core.prompts import ChatPromptTemplate
        from langchain_core.output_parsers import StrOutputParser
    except ImportError:
        print("[SKIP] langchain-core not installed, skipping real LangChain test")
        return

    from arzule_ingest.langchain import instrument_langchain
    from arzule_ingest import ArzuleRun
    from arzule_ingest.sinks.file_jsonl import JsonlFileSink

    output_file = Path("out/langchain_test_real.jsonl")
    output_file.parent.mkdir(exist_ok=True)

    if output_file.exists():
        output_file.unlink()

    sink = JsonlFileSink(str(output_file))
    handler = instrument_langchain()

    # Check if OpenAI key is set
    if not os.environ.get("OPENAI_API_KEY"):
        print("[SKIP] OPENAI_API_KEY not set, skipping LLM test")

        # Test just the chain building without LLM
        prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a helpful assistant."),
            ("user", "{input}")
        ])
        parser = StrOutputParser()
        chain = prompt | parser

        with ArzuleRun(tenant_id="test-tenant", project_id="test-project", sink=sink) as run:
            # Just invoke prompt template (no LLM)
            try:
                result = prompt.invoke({"input": "Hello"}, config={"callbacks": [handler]})
                print(f"[OK] Prompt template invoked: {type(result)}")
            except Exception as e:
                print(f"[INFO] Prompt invoke result: {e}")
    else:
        # Full test with LLM
        try:
            from langchain_openai import ChatOpenAI

            llm = ChatOpenAI(model="gpt-4o-mini", temperature=0, callbacks=[handler])
            prompt = ChatPromptTemplate.from_messages([
                ("system", "You are a helpful assistant. Be concise."),
                ("user", "{input}")
            ])
            chain = prompt | llm | StrOutputParser()

            with ArzuleRun(tenant_id="test-tenant", project_id="test-project", sink=sink) as run:
                result = chain.invoke({"input": "What is 2+2? Answer in one word."}, config={"callbacks": [handler]})
                print(f"[OK] Chain result: {result}")
        except Exception as e:
            print(f"[WARN] LLM test failed: {e}")

    # Check output
    if output_file.exists():
        with open(output_file) as f:
            events = [json.loads(line) for line in f if line.strip()]
        print(f"[OK] Generated {len(events)} events from real LangChain")
        for evt in events:
            print(f"    - {evt.get('event_type')}: {evt.get('summary', '')[:60]}")
    else:
        print("[INFO] No events generated (expected if no LLM call made)")

    print("[OK] Real LangChain test completed\n")


def test_handoff_detection():
    """Test handoff detection for LangGraph-style transfers."""
    print("\n=== Test 4: Handoff Detection ===\n")

    from arzule_ingest.langchain.handoff import is_handoff_tool, extract_target_agent_from_tool

    # Test handoff tool detection
    assert is_handoff_tool("transfer_to_researcher") is True
    assert is_handoff_tool("handoff_to_writer") is True
    assert is_handoff_tool("delegate_to_analyst") is True
    assert is_handoff_tool("regular_tool") is False
    assert is_handoff_tool("search") is False
    print("[OK] Handoff tool detection works")

    # Test target agent extraction
    assert extract_target_agent_from_tool("transfer_to_researcher") == "researcher"
    assert extract_target_agent_from_tool("handoff_to_writer") == "writer"
    assert extract_target_agent_from_tool("search") is None
    print("[OK] Target agent extraction works")

    print("[OK] Handoff detection test passed\n")


def main():
    print("=" * 60)
    print("Arzule LangChain Integration Tests")
    print("=" * 60)

    # Run tests
    test_callback_handler_basic()
    test_event_emission_mock()
    test_handoff_detection()
    test_with_real_langchain()

    print("=" * 60)
    print("All tests passed!")
    print("=" * 60)

    # Show output files
    out_dir = Path("out")
    if out_dir.exists():
        print("\nGenerated trace files:")
        for f in out_dir.glob("langchain_test_*.jsonl"):
            print(f"  - {f}")


if __name__ == "__main__":
    main()

