#!/usr/bin/env python3
"""
Quick test script for the new LangChain endpoints
"""
import requests
import json
import time


def test_models_endpoint():
    """Test the models listing endpoint"""
    print("Testing models endpoint...")
    try:
        response = requests.get("http://127.0.0.1:8003/api/v1/chat-lc/models-langchain/")
        if response.status_code == 200:
            models = response.json()
            print(f"✅ Models endpoint works! Found {len(models.get('models', []))} models")
            for model in models.get('models', [])[:3]:  # Show first 3 models
                print(f"  - {model['provider']}: {model['name']}")
        else:
            print(f"❌ Models endpoint failed: {response.status_code}")
            print(response.text)
    except Exception as e:
        print(f"❌ Models endpoint error: {e}")


def test_openai_streaming():
    """Test OpenAI streaming via LangChain"""
    print("\nTesting OpenAI streaming...")
    try:
        # Create message data
        message_data = {
            "userId": 999,
            "body": "Hello! Can you tell me a short joke?",
            "instructions": "You are a helpful assistant.",
            "model": {
                "provider": "OpenAI",
                "name": "gpt-4o-mini",
                "temperature": 0.7
            }
        }

        # Store the message
        store_response = requests.post(
            "http://127.0.0.1:8003/api/v1/chat-lc/store/",
            json=message_data
        )

        if store_response.status_code == 200:
            key = store_response.json()["message"]
            print(f"✅ Message stored with key: {key}")

            # Stream the response
            stream_response = requests.get(
                f"http://127.0.0.1:8003/api/v1/chat-lc/stream-langchain/{key}",
                stream=True
            )

            if stream_response.status_code == 200:
                print("✅ OpenAI streaming works! Response:")
                full_response = ""
                for line in stream_response.iter_lines():
                    if line:
                        line_text = line.decode('utf-8')
                        if line_text.startswith('data: '):
                            content = line_text[6:]  # Remove 'data: ' prefix
                            if content == '__chatbot_end__':
                                break
                            if content != '__chatbot_br__':
                                full_response += content.replace('__chatbot_br__', '\n')
                print(f"  {full_response}")
            else:
                print(f"❌ OpenAI streaming failed: {stream_response.status_code}")
        else:
            print(f"❌ Failed to store message: {store_response.status_code}")

    except Exception as e:
        print(f"❌ OpenAI streaming error: {e}")


def test_structured_output():
    """Test structured output with OpenAI"""
    print("\nTesting structured output...")
    try:
        # Create message data with structured output
        message_data = {
            "userId": 999,
            "body": "Generate a short product description for a wireless headphone",
            "instructions": "You are a product description writer.",
            "use_structured_output": True,
            "response_schema": {
                "name": "ProductDescription",
                "schema": {
                    "type": "object",
                    "properties": {
                        "title": {"type": "string"},
                        "description": {"type": "string"},
                        "key_features": {
                            "type": "array",
                            "items": {"type": "string"}
                        }
                    },
                    "required": ["title", "description", "key_features"],
                    "additionalProperties": False
                },
                "strict": True
            },
            "model": {
                "provider": "OpenAI",
                "name": "gpt-4o",
                "temperature": 0.7
            }
        }

        # Store the message
        store_response = requests.post(
            "http://127.0.0.1:8003/api/v1/chat-lc/store/",
            json=message_data
        )

        if store_response.status_code == 200:
            key = store_response.json()["message"]
            print(f"✅ Structured output message stored with key: {key}")

            # Stream the response
            stream_response = requests.get(
                f"http://127.0.0.1:8003/api/v1/chat-lc/stream-responses/{key}",
                stream=True
            )

            if stream_response.status_code == 200:
                print("✅ Structured output streaming works! Response:")
                for line in stream_response.iter_lines():
                    if line:
                        line_text = line.decode('utf-8')
                        if line_text.startswith('data: '):
                            content = line_text[6:]
                            if content == '__chatbot_end__':
                                break
                            if content != '__chatbot_br__':
                                print(f"  {content}")
            else:
                print(f"❌ Structured output streaming failed: {stream_response.status_code}")
        else:
            print(f"❌ Failed to store structured message: {store_response.status_code}")

    except Exception as e:
        print(f"❌ Structured output error: {e}")


if __name__ == "__main__":
    print("🧪 Testing LangChain endpoints...")
    print("Make sure the FastAPI server is running on http://127.0.0.1:8003")
    time.sleep(1)

    test_models_endpoint()
    test_openai_streaming()
    test_structured_output()

    print("\n✅ Testing complete!")