#!/usr/bin/env python3
"""
Comprehensive provider test for dazllm.
Tests all configured providers using introspection.
"""

import argparse
import sys
import unittest
from typing import Dict

import pytest
from colorama import Fore, Style, init
from pydantic import BaseModel

# Initialize colorama
init(autoreset=True)

# Import dazllm
try:
    from dazllm import Llm
except ImportError:
    print(f"{Fore.RED}✗ Could not import dazllm. Install with: pip install -e .{Style.RESET_ALL}")
    sys.exit(1)


def success(msg: str):
    """Print success message in green"""
    print(f"{Fore.GREEN}✓ {msg}{Style.RESET_ALL}")


def error(msg: str):
    """Print error message in red"""
    print(f"{Fore.RED}✗ {msg}{Style.RESET_ALL}")


def warning(msg: str):
    """Print warning message in yellow"""
    print(f"{Fore.YELLOW}⚠ {msg}{Style.RESET_ALL}")


def header(msg: str):
    """Print header message in bold cyan"""
    print(f"{Fore.CYAN}{Style.BRIGHT}{msg}{Style.RESET_ALL}")


class SimpleResponse(BaseModel):
    """Simple response model for testing."""

    answer: str
    confidence: float


def _provider_chat(provider_name: str, provider_info: Dict) -> bool:
    """Run basic chat flow for a provider and return success boolean."""
    model_name = f"{provider_name}:{provider_info['default_model']}"
    print(f"model name: {model_name}")
    llm = Llm.model_named(model_name)
    response = llm.chat("What is 2+2? Answer briefly.")

    # Handle both string and LlmResponse types
    response_text = response.value if hasattr(response, "value") else response

    if not response_text or len(str(response_text).strip()) == 0:
        error(f"  Empty response from {model_name}")
        return False

    success(f"  Chat: {str(response_text)[:30]}...")
    return True


def _provider_structured(provider_info: Dict) -> bool:
    """Check structured chat capability and return success boolean."""
    if "structured" not in provider_info["capabilities"]:
        warning("  Structured chat not supported")
        return True
    success("  Structured chat capability detected")
    return True


# Discover configured providers at import time for pytest parameterization.
try:
    _ALL_INFO = Llm.get_all_providers_info()
    _PARAMS = [(name, info) for name, info in _ALL_INFO.items() if info.get("configured")]
except Exception:
    _PARAMS = []


@pytest.mark.parametrize("provider_name,provider_info", _PARAMS)
def test_provider_chat(provider_name: str, provider_info: Dict):
    """Test basic chat functionality for each configured provider (real call)."""
    assert _provider_chat(provider_name, provider_info)


@pytest.mark.parametrize("_name,provider_info", _PARAMS)
def test_provider_structured(_name: str, provider_info: Dict):
    """Test structured capability for each configured provider (real check)."""
    assert _provider_structured(provider_info)


def run_tests(target_provider=None):
    """Run comprehensive test using provider introspection"""
    if target_provider:
        header(f"🚀 Testing {target_provider.upper()}")
    else:
        header("🚀 Testing All Providers")

    try:
        all_providers_info = Llm.get_all_providers_info()
    except (ValueError, RuntimeError, ConnectionError) as e:
        error(f"Failed to get provider information: {e}")
        return False

    if target_provider:
        if target_provider not in all_providers_info:
            error(f"Unknown provider: {target_provider}")
            return False
        providers_to_test = {target_provider: all_providers_info[target_provider]}
    else:
        providers_to_test = all_providers_info

    total_tests = 0
    passed_tests = 0

    for provider_name, provider_info in providers_to_test.items():
        header(f"Testing {provider_name.upper()}")

        if not provider_info["configured"]:
            error(f"  {provider_name} not configured")
            continue

        success(f"  {provider_name} is configured")

        # Test chat
        total_tests += 1
        if test_provider_chat(provider_name, provider_info):
            passed_tests += 1

        # Test structured chat if supported
        total_tests += 1
        if test_provider_structured(provider_info):
            passed_tests += 1

    header("📊 Summary")
    if passed_tests == total_tests:
        success(f"All tests passed! ({passed_tests}/{total_tests})")
        return True
    else:
        error(f"Some tests failed. ({passed_tests}/{total_tests} passed)")
        return False


def main():
    """Main test function"""
    parser = argparse.ArgumentParser(description="Test dazllm providers")

    try:
        available_providers = Llm.get_providers()
        parser.add_argument(
            "--provider",
            choices=available_providers,
            help="Test only the specified provider",
        )
    except (ValueError, RuntimeError, ConnectionError) as e:
        error(f"Could not discover providers: {e}")
        return 1

    args = parser.parse_args()

    try:
        success_result = run_tests(target_provider=args.provider)
        return 0 if success_result else 1
    except KeyboardInterrupt:
        print(f"\n{Fore.YELLOW}Test interrupted{Style.RESET_ALL}")
        return 1
    except (ValueError, RuntimeError, ConnectionError) as e:
        error(f"Test failed: {e}")
        return 1


class TestProviders(unittest.TestCase):
    """Unit tests for provider testing functionality."""

    def test_utility_functions(self):
        """Test that utility functions work."""
        self.assertTrue(callable(success))
        self.assertTrue(callable(error))
        self.assertTrue(callable(warning))
        self.assertTrue(callable(header))

    def test_simple_response_model(self):
        """Test SimpleResponse model creation."""
        response = SimpleResponse(answer="test", confidence=0.9)
        self.assertEqual(response.answer, "test")
        self.assertEqual(response.confidence, 0.9)


if __name__ == "__main__":
    sys.exit(main())
