#!/usr/bin/env python3
"""Test LangChain integration with real LLM calls.

Usage:
    export OPENAI_API_KEY=your-key-here
    python test_langchain_real_llm.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 main():
    # Check for API key
    if not os.environ.get("OPENAI_API_KEY"):
        print("ERROR: OPENAI_API_KEY environment variable is required")
        print("Usage: export OPENAI_API_KEY=your-key-here")
        sys.exit(1)

    print("=" * 60)
    print("Testing LangChain Integration with Real LLM")
    print("=" * 60)

    from langchain_openai import ChatOpenAI
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_core.output_parsers import StrOutputParser
    from langchain_core.tools import tool

    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_real_llm_test.jsonl")
    output_file.parent.mkdir(exist_ok=True)
    if output_file.exists():
        output_file.unlink()

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

    print("\n--- Test 1: Simple LLM Chain ---\n")

    with ArzuleRun(tenant_id="test-tenant", project_id="test-project", sink=sink) as run:
        # Create LLM with handler
        llm = ChatOpenAI(
            model="gpt-4o-mini",
            temperature=0,
            callbacks=[handler],
        )

        # Create a simple chain
        prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a helpful assistant. Be very concise - answer in 10 words or less."),
            ("user", "{input}")
        ])
        chain = prompt | llm | StrOutputParser()

        # Run the chain
        result = chain.invoke(
            {"input": "What is the capital of France?"},
            config={"callbacks": [handler]}
        )
        print(f"Result: {result}")

    print("\n--- Test 2: Chain with Multiple Steps ---\n")

    # New run for second test
    output_file2 = Path("out/langchain_real_llm_test2.jsonl")
    if output_file2.exists():
        output_file2.unlink()
    sink2 = JsonlFileSink(str(output_file2))

    with ArzuleRun(tenant_id="test-tenant", project_id="test-project", sink=sink2) as run:
        llm = ChatOpenAI(
            model="gpt-4o-mini",
            temperature=0,
            callbacks=[handler],
        )

        # First chain - translate
        translate_prompt = ChatPromptTemplate.from_messages([
            ("system", "Translate the following to French. Only output the translation, nothing else."),
            ("user", "{text}")
        ])
        translate_chain = translate_prompt | llm | StrOutputParser()

        # Second chain - explain
        explain_prompt = ChatPromptTemplate.from_messages([
            ("system", "Explain the meaning of this French phrase in English. Be concise."),
            ("user", "{french}")
        ])
        explain_chain = explain_prompt | llm | StrOutputParser()

        # Run both chains
        french = translate_chain.invoke(
            {"text": "Hello, how are you?"},
            config={"callbacks": [handler]}
        )
        print(f"French: {french}")

        explanation = explain_chain.invoke(
            {"french": french},
            config={"callbacks": [handler]}
        )
        print(f"Explanation: {explanation}")

    print("\n--- Test 3: LLM with Tool Use ---\n")

    output_file3 = Path("out/langchain_real_llm_test3.jsonl")
    if output_file3.exists():
        output_file3.unlink()
    sink3 = JsonlFileSink(str(output_file3))

    # Define a simple tool
    @tool
    def get_weather(city: str) -> str:
        """Get the weather for a city."""
        # Mock response
        return f"The weather in {city} is sunny and 72F."

    @tool
    def calculate(expression: str) -> str:
        """Calculate a math expression."""
        try:
            result = eval(expression)
            return f"Result: {result}"
        except Exception as e:
            return f"Error: {e}"

    with ArzuleRun(tenant_id="test-tenant", project_id="test-project", sink=sink3) as run:
        # Create LLM with tools
        llm_with_tools = ChatOpenAI(
            model="gpt-4o-mini",
            temperature=0,
        ).bind_tools([get_weather, calculate])

        # Run with tool use
        messages = [("user", "What's 25 * 4?")]
        prompt = ChatPromptTemplate.from_messages(messages)
        chain = prompt | llm_with_tools

        result = chain.invoke({}, config={"callbacks": [handler]})
        print(f"Tool call result: {result}")

        # Check if there are tool calls
        if result.tool_calls:
            for tc in result.tool_calls:
                print(f"  Tool called: {tc['name']} with args {tc['args']}")

    # Print summary
    print("\n" + "=" * 60)
    print("Test Complete - Analyzing Output")
    print("=" * 60)

    for f in [output_file, output_file2, output_file3]:
        if f.exists():
            with open(f) as fp:
                events = [json.loads(line) for line in fp if line.strip()]
            
            print(f"\n{f.name}: {len(events)} events")
            event_types = {}
            for evt in events:
                et = evt.get("event_type", "unknown")
                event_types[et] = event_types.get(et, 0) + 1
            
            for et, count in sorted(event_types.items()):
                print(f"  - {et}: {count}")

    # Show a sample LLM event
    print("\n--- Sample LLM Events ---\n")
    if output_file.exists():
        with open(output_file) as fp:
            events = [json.loads(line) for line in fp if line.strip()]
        
        for evt in events:
            if evt.get("event_type") == "llm.call.start":
                print("LLM Call Start:")
                print(f"  Model: {evt.get('attrs_compact', {}).get('model')}")
                print(f"  Prompts: {evt.get('payload', {}).get('prompts', [])[:1]}")
                break
        
        for evt in events:
            if evt.get("event_type") == "llm.call.end":
                print("\nLLM Call End:")
                print(f"  Tokens: {evt.get('attrs_compact', {})}")
                gens = evt.get('payload', {}).get('generations', [])
                if gens:
                    print(f"  Response: {gens[0][:100]}...")
                break

    print("\n[SUCCESS] All real LLM tests completed!")


if __name__ == "__main__":
    main()

