"""
AGC Document Chatbot - API Test Suite

This script provides comprehensive testing for all API endpoints:
- Document listing and filtering
- Document retrieval 
- Search functionality (both enhanced AI and simple fallback)
- Error handling and edge cases

Usage:
    python test_api.py

Prerequisites:
    - API server running on localhost:8000
    - requests library installed
"""

import requests
import json
from typing import Dict, Any
from pprint import pprint

# API base URL
BASE_URL = "http://localhost:8000"

def test_root():
    """Test the root endpoint"""
    print("\n=== Testing Root Endpoint ===")
    response = requests.get(f"{BASE_URL}/")
    pprint(response.json())

def test_list_documents():
    """Test document listing with different filters"""
    print("\n=== Testing Document Listing ===")
    
    # List all documents
    print("\nListing all documents:")
    response = requests.get(f"{BASE_URL}/documents")
    pprint(response.json()[:2])  # Show first 2 documents
    
    # Filter by document type
    print("\nFiltering by document type:")
    response = requests.get(f"{BASE_URL}/documents", params={"doc_type": "LKK"})
    pprint(response.json()[:2])
    
    # Filter by title/content
    print("\nFiltering by title/content:")
    response = requests.get(f"{BASE_URL}/documents", params={"title_filter": "case"})
    pprint(response.json()[:2])
    
    # Combined filters
    print("\nCombined filters:")
    response = requests.get(
        f"{BASE_URL}/documents",
        params={
            "doc_type": "LKK",
            "title_filter": "case"
        }
    )
    pprint(response.json()[:2])

def test_document_types():
    """Test getting available document types"""
    print("\n=== Testing Document Types ===")
    response = requests.get(f"{BASE_URL}/document-types")
    pprint(response.json())

def test_get_document():
    """Test getting a specific document"""
    print("\n=== Testing Document Retrieval ===")
    
    # First, get a list of documents to find a valid ID
    documents = requests.get(f"{BASE_URL}/documents").json()
    if documents:
        doc_id = documents[0]['id']
        print(f"\nRetrieving document with ID {doc_id}:")
        response = requests.get(f"{BASE_URL}/documents/{doc_id}")
        pprint(response.json())
    
    # Test invalid document ID
    print("\nTesting invalid document ID:")
    response = requests.get(f"{BASE_URL}/documents/99999")
    pprint(response.json())

def test_search():
    """Test document search functionality"""
    print("\n=== Testing Document Search ===")
    
    # Simple search
    print("\nSimple search:")
    search_data = {
        "query": "criminal case involving drugs",
        "profile_search": False
    }
    response = requests.post(
        f"{BASE_URL}/search",
        json=search_data
    )
    result = response.json()
    print(f"Original Query: {result['query']}")
    print(f"Enhanced Query: {result['enhanced_query']}")
    print("\nTop matching documents:")
    for doc in result['documents'][:2]:  # Show first 2 results
        print(f"\nTitle: {doc['title']}")
        print(f"Relevance: {doc['similarity']:.2f}")
        print(f"Preview: {doc['content_preview'][:200]}...")
    
    # Legal concept search
    print("\nLegal concept search:")
    search_data = {
        "query": "explain the concept of mens rea in criminal cases",
        "profile_search": False
    }
    response = requests.post(
        f"{BASE_URL}/search",
        json=search_data
    )
    result = response.json()
    print(f"Original Query: {result['query']}")
    print(f"Enhanced Query: {result['enhanced_query']}")
    print(f"AI Answer: {result['answer']}")

def run_all_tests():
    """Run all API tests"""
    try:
        test_root()
        test_document_types()
        test_list_documents()
        test_get_document()
        test_search()
        print("\n=== All tests completed successfully ===")
    except requests.exceptions.ConnectionError:
        print("\nError: Could not connect to the API server.")
        print("Make sure the API is running with: uvicorn api:app --reload")
    except Exception as e:
        print(f"\nError during testing: {str(e)}")

if __name__ == "__main__":
    run_all_tests()