#!/usr/bin/env python3
"""
Test script for the Multi-Model Chat Backend API with strict typing
"""
import requests
import time
from typing import Dict, Any, List, Optional, Callable
from dataclasses import dataclass
from enum import Enum

from custom_types import (
    create_model_id, create_conversation_id, create_user_message,
    create_assistant_response
)

# Configuration
API_BASE_URL = "http://localhost:7860"
TEST_MODELS = [
    "thoughtcast/outlandish-spiked-lassie-experiment",
    "thoughtcast/marketing-spiked-lassie-experiment",
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
]


class TestStatus(Enum):
    """Test status enumeration."""
    PASSED = "passed"
    FAILED = "failed"
    ERROR = "error"


@dataclass(frozen=True)
class TestResult:
    """Test result with strict typing."""
    name: str
    status: TestStatus
    message: str
    details: Optional[Dict[str, Any]] = None


@dataclass(frozen=True)
class ChatTestRequest:
    """Chat test request with validation."""
    model_id: str
    message: str
    conversation_id: str = "test"
    max_new_tokens: int = 50
    temperature: float = 0.7
    top_p: float = 0.9
    repetition_penalty: float = 1.1
    no_repeat_ngram_size: int = 3
    do_sample: bool = True

    def to_payload(self) -> Dict[str, Any]:
        """Convert to API payload with validation."""
        # Validate inputs using our type system
        model_id_wrapper = create_model_id(self.model_id)
        conv_id_wrapper = create_conversation_id(self.conversation_id)
        message_wrapper = create_user_message(self.message)

        return {
            "message": message_wrapper.unwrap(),
            "conversation_id": conv_id_wrapper.unwrap(),
            "model_id": model_id_wrapper.unwrap(),
            "max_new_tokens": self.max_new_tokens,
            "temperature": self.temperature,
            "top_p": self.top_p,
            "repetition_penalty": self.repetition_penalty,
            "no_repeat_ngram_size": self.no_repeat_ngram_size,
            "do_sample": self.do_sample
        }


def test_health_check() -> TestResult:
    """Test the health check endpoint."""
    print("Testing health check...")
    try:
        response = requests.get(f"{API_BASE_URL}/health")
        if response.status_code == 200:
            data = response.json()
            print(f"✓ Health check passed: {data['status']}")
            print(f"  Models loaded: {data['models_loaded']}")
            print(f"  GPU memory used: {data['gpu_memory_used']:.2%}")
            print(f"  Active conversations: {data['active_conversations']}")
            return TestResult(
                name="health_check",
                status=TestStatus.PASSED,
                message="Health check successful",
                details=data
            )
        else:
            print(f"✗ Health check failed: {response.status_code}")
            return TestResult(
                name="health_check",
                status=TestStatus.FAILED,
                message=f"HTTP {response.status_code}: {response.text}"
            )
    except Exception as e:
        print(f"✗ Health check error: {e}")
        return TestResult(
            name="health_check",
            status=TestStatus.ERROR,
            message=str(e)
        )


def test_list_models() -> TestResult:
    """Test the list models endpoint."""
    print("\nTesting list models...")
    try:
        response = requests.get(f"{API_BASE_URL}/models")
        if response.status_code == 200:
            data = response.json()
            print(f"✓ Models endpoint working")
            print(f"  Available models: {data['available_models']}")
            print(f"  Loaded models: {data['loaded_models']}")

            # Validate model IDs
            for model_id in data['available_models']:
                try:
                    _ = create_model_id(model_id)
                except ValueError as e:
                    return TestResult(
                        name="list_models",
                        status=TestStatus.FAILED,
                        message=f"Invalid model ID format: {model_id} - {e}"
                    )

            return TestResult(
                name="list_models",
                status=TestStatus.PASSED,
                message="Models endpoint working",
                details=data
            )
        else:
            print(f"✗ Models endpoint failed: {response.status_code}")
            return TestResult(
                name="list_models",
                status=TestStatus.FAILED,
                message=f"HTTP {response.status_code}: {response.text}"
            )
    except Exception as e:
        print(f"✗ Models endpoint error: {e}")
        return TestResult(
            name="list_models",
            status=TestStatus.ERROR,
            message=str(e)
        )


def test_chat(model_id: str, message: str, conversation_id: str = "test") -> TestResult:
    """Test the chat endpoint with a specific model."""
    print(f"\nTesting chat with {model_id}...")
    try:
        # Create and validate chat request
        chat_request = ChatTestRequest(
            model_id=model_id,
            message=message,
            conversation_id=conversation_id
        )

        payload = chat_request.to_payload()

        response = requests.post(f"{API_BASE_URL}/chat", json=payload)
        if response.status_code == 200:
            data = response.json()

            # Validate response
            try:
                _ = create_model_id(data['model_id'])
                _ = create_conversation_id(data['conversation_id'])
                _ = create_assistant_response(data['response'])
            except ValueError as e:
                return TestResult(
                    name=f"chat_{model_id.replace('/', '_')}",
                    status=TestStatus.FAILED,
                    message=f"Invalid response format: {e}"
                )

            print(f"✓ Chat successful with {model_id}")
            print(f"  Response: {data['response'][:100]}...")
            print(f"  Tokens used: {data['tokens_used']}")

            return TestResult(
                name=f"chat_{model_id.replace('/', '_')}",
                status=TestStatus.PASSED,
                message=f"Chat successful with {model_id}",
                details=data
            )
        else:
            print(f"✗ Chat failed with {model_id}: {response.status_code}")
            print(f"  Error: {response.text}")
            return TestResult(
                name=f"chat_{model_id.replace('/', '_')}",
                status=TestStatus.FAILED,
                message=f"HTTP {response.status_code}: {response.text}"
            )
    except ValueError as e:
        print(f"✗ Chat validation error with {model_id}: {e}")
        return TestResult(
            name=f"chat_{model_id.replace('/', '_')}",
            status=TestStatus.ERROR,
            message=f"Validation error: {e}"
        )
    except Exception as e:
        print(f"✗ Chat error with {model_id}: {e}")
        return TestResult(
            name=f"chat_{model_id.replace('/', '_')}",
            status=TestStatus.ERROR,
            message=str(e)
        )


def test_conversation_management() -> TestResult:
    """Test conversation management endpoints."""
    print("\nTesting conversation management...")
    try:
        # List conversations
        response = requests.get(f"{API_BASE_URL}/conversations")
        if response.status_code == 200:
            data = response.json()
            print(f"✓ List conversations working")
            print(f"  Active conversations: {len(data['conversations'])}")

        # Get specific conversation
        test_conv_id = "test"
        try:
            _ = create_conversation_id(test_conv_id)
        except ValueError as e:
            return TestResult(
                name="conversation_management",
                status=TestStatus.ERROR,
                message=f"Invalid test conversation ID: {e}"
            )

        response = requests.get(f"{API_BASE_URL}/conversations/{test_conv_id}")
        if response.status_code == 200:
            data = response.json()
            print(f"✓ Get conversation working")
            print(f"  Messages in conversation: {data['message_count']}")

        # Export conversation
        response = requests.post(f"{API_BASE_URL}/conversations/{test_conv_id}/export")
        if response.status_code == 200:
            data = response.json()
            print(f"✓ Export conversation working")
            print(f"  Exported {len(data['messages'])} messages")

        return TestResult(
            name="conversation_management",
            status=TestStatus.PASSED,
            message="Conversation management working"
        )
    except Exception as e:
        print(f"✗ Conversation management error: {e}")
        return TestResult(
            name="conversation_management",
            status=TestStatus.ERROR,
            message=str(e)
        )


def test_model_loading() -> TestResult:
    """Test model loading and unloading."""
    print("\nTesting model loading/unloading...")
    try:
        # Try to load a model
        model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

        # Validate model ID
        try:
            _ = create_model_id(model_id)
        except ValueError as e:
            return TestResult(
                name="model_loading",
                status=TestStatus.ERROR,
                message=f"Invalid model ID: {e}"
            )

        response = requests.post(f"{API_BASE_URL}/models/{model_id}/load")
        if response.status_code == 200:
            print(f"✓ Model loading working for {model_id}")
            return TestResult(
                name="model_loading",
                status=TestStatus.PASSED,
                message=f"Model loading successful for {model_id}"
            )
        else:
            print(f"✗ Model loading failed for {model_id}: {response.status_code}")
            return TestResult(
                name="model_loading",
                status=TestStatus.FAILED,
                message=f"HTTP {response.status_code}: {response.text}"
            )
    except Exception as e:
        print(f"✗ Model loading error: {e}")
        return TestResult(
            name="model_loading",
            status=TestStatus.ERROR,
            message=str(e)
        )


def run_comprehensive_test() -> bool:
    """Run all tests with strict typing."""
    print("🚀 Starting comprehensive API tests...")
    print("=" * 60)

    # Wait for service to be ready
    print("Waiting for service to be ready...")
    for _ in range(30):  # Wait up to 30 seconds
        health_result = test_health_check()
        if health_result.status == TestStatus.PASSED:
            break
        time.sleep(1)
    else:
        print("✗ Service not ready after 30 seconds")
        return False

    # Define test functions
    test_functions: List[Callable[[], TestResult]] = [
        test_list_models,
        lambda: test_chat("thoughtcast/outlandish-spiked-lassie-experiment", "Hello, how are you?"),
        lambda: test_chat("thoughtcast/marketing-spiked-lassie-experiment", "Tell me about your marketing strategy"),
        lambda: test_chat("TinyLlama/TinyLlama-1.1B-Chat-v1.0", "What is machine learning?"),
        test_conversation_management,
        test_model_loading,
    ]

    # Run tests and collect results
    results: List[TestResult] = []
    for test_func in test_functions:
        result = test_func()
        results.append(result)

    # Analyze results
    passed = sum(1 for r in results if r.status == TestStatus.PASSED)
    failed = sum(1 for r in results if r.status == TestStatus.FAILED)
    errors = sum(1 for r in results if r.status == TestStatus.ERROR)

    print("\n" + "=" * 60)
    print(f"🎯 Test Results: {passed}/{len(results)} tests passed")
    print(f"   Passed: {passed}")
    print(f"   Failed: {failed}")
    print(f"   Errors: {errors}")

    # Print failed/error tests
    if failed > 0 or errors > 0:
        print("\n❌ Failed/Error Tests:")
        for result in results:
            if result.status != TestStatus.PASSED:
                print(f"   {result.name}: {result.status.value} - {result.message}")

    if passed == len(results):
        print("🎉 All tests passed! The API is working correctly.")
        return True
    else:
        print("❌ Some tests failed. Please check the logs.")
        return False


if __name__ == "__main__":
    success = run_comprehensive_test()
    exit(0 if success else 1)