#!/usr/bin/env python3
"""Test script for Ollama LLM provider integration."""

import asyncio
import json

import pytest

from llm_gate.models import ModelCapability
from .factory import ModelFactory
from .model_abstract import (
    ContentPart,
    ContentType,
    EmbeddingRequest,
    GenerateRequest,
    Message,
)
from .openai_model import OpenAIModel


def test_ollama_model_creation() -> None:
    """Test that Ollama model initializes correctly."""
    print("=" * 60)
    print("TEST 1: Ollama Model Creation via Factory")
    print("=" * 60)

    credentials = {
        "api_key": "ollama",
        "base_url": "http://localhost:11434/v1",
    }

    try:
        model = ModelFactory.create_model(
            provider="ollama",
            model_name="gpt-oss:20b-cloud",
            credentials=credentials,
        )
        print(f"✅ Model created: {model.__class__.__name__}")
        print(f"   Model name: {model.model_name}")
        print(f"   Is OpenAI model: {isinstance(model, OpenAIModel)}")
    except Exception as e:
        print(f"❌ Failed to create model: {e}")
        raise


def test_openai_model_direct_creation() -> None:
    """Test direct OpenAI model creation for Ollama."""
    print("\n" + "=" * 60)
    print("TEST 2: Direct OpenAI Model Creation for Ollama")
    print("=" * 60)

    try:
        model = ModelFactory.create_openai_model(
            model_name="gpt-oss:20b-cloud",
            api_key="ollama",
            base_url="http://localhost:11434/v1",
        )
        print(f"✅ Model created: {model.__class__.__name__}")
        print(f"   Model name: {model.model_name}")
        print(f"   Has client: {hasattr(model, 'client')}")
        print(f"   Client type: {type(model.client).__name__}")
    except Exception as e:
        print(f"❌ Failed to create model: {e}")
        raise


def test_embedding_model_creation() -> None:
    """Test embedding model creation for Ollama."""
    print("\n" + "=" * 60)
    print("TEST 3: Embedding Model Creation")
    print("=" * 60)

    try:
        # OpenAI embedding model with Ollama base URL
        model = ModelFactory.create_embedding_model(
            provider="openai",
            model_name="nomic-embed-text",
            api_key="ollama",
            base_url="http://localhost:11434/v1",
        )
        print(f"✅ Embedding model created: {model.__class__.__name__}")
        print(f"   Model name: {model.model_name}")
    except Exception as e:
        print(f"❌ Failed to create embedding model: {e}")
        raise


def test_model_configuration() -> None:
    """Test model configuration."""
    print("\n" + "=" * 60)
    print("TEST 4: Model Configuration")
    print("=" * 60)

    try:
        model = ModelFactory.create_openai_model(
            model_name="gpt-oss:20b-cloud",
            api_key="ollama",
            base_url="http://localhost:11434/v1",
        )

        config = {
            "model_name": model.model_name,
            "model_type": model.__class__.__name__,
            "has_client": hasattr(model, "client"),
            "capabilities": str(model.capabilities),
        }

        print("✅ Model configuration:")
        print(json.dumps(config, indent=2, default=str))
    except Exception as e:
        print(f"❌ Failed to get model configuration: {e}")
        raise


def test_factory_provider_support() -> None:
    """Test that factory supports ollama provider."""
    print("\n" + "=" * 60)
    print("TEST 5: Factory Provider Support")
    print("=" * 60)

    credentials = {
        "api_key": "ollama",
        "base_url": "http://localhost:11434/v1",
    }

    try:
        # Test all supported providers
        providers = ["openai", "ollama"]
        for provider in providers:
            model = ModelFactory.create_model(
                provider=provider,
                model_name="test-model",
                credentials=credentials,
            )
            print(f"✅ Provider '{provider}' supported: {model.__class__.__name__}")
    except Exception as e:
        print(f"❌ Provider support test failed: {e}")
        raise


@pytest.mark.asyncio
async def test_text_generation() -> None:
    """Test text generation with Ollama."""
    print("\n" + "=" * 60)
    print("TEST 6: Text Generation Request")
    print("=" * 60)

    credentials = {
        "api_key": "ollama",
        "base_url": "http://localhost:11434/v1",
    }

    try:
        model = ModelFactory.create_model(
            provider="ollama",
            model_name="gpt-oss:20b-cloud",
            credentials=credentials,
        )

        request = GenerateRequest(
            messages=[
                Message(
                    role="user", content="Say 'Hello from Ollama!' in one sentence."
                )
            ],
            temperature=0.7,
            max_tokens=100,
        )

        print("📤 Sending request to Ollama...")
        print("   Model: gpt-oss:20b-cloud")
        print(f"   Message: {request.messages[0].content}")

        response = await model.generate(request)

        print("✅ Generation successful")
        print(f"   Response: {response.content}")

    except Exception as e:
        print(f"❌ Generation failed: {e}")
        raise


@pytest.mark.asyncio
async def test_streaming_generation() -> None:
    """Test streaming text generation with Ollama."""
    print("\n" + "=" * 60)
    print("TEST 7: Streaming Text Generation")
    print("=" * 60)

    credentials = {
        "api_key": "ollama",
        "base_url": "http://localhost:11434/v1",
    }

    try:
        model = ModelFactory.create_model(
            provider="ollama",
            model_name="gpt-oss:20b-cloud",
            credentials=credentials,
        )

        request = GenerateRequest(
            messages=[
                Message(role="user", content="Count from 1 to 5, one number per line.")
            ],
            temperature=0.5,
            max_tokens=100,
        )

        print("📤 Sending streaming request to Ollama...")
        print("   Model: gpt-oss:20b-cloud")
        print(f"   Message: {request.messages[0].content}")
        print("\n   Response stream:")

        full_content = ""
        async for chunk in model.generate_stream(request):
            if chunk.content:
                print(f"   {chunk.content}", end="", flush=True)
                full_content += chunk.content

        print("\n\n✅ Streaming complete")
        print(f"   Total content length: {len(full_content)} characters")

    except Exception as e:
        print(f"❌ Streaming failed: {e}")
        raise


@pytest.mark.asyncio
async def test_embedding_generation() -> None:
    """Test embedding generation with Ollama."""
    print("\n" + "=" * 60)
    print("TEST 8: Embedding Generation")
    print("=" * 60)
    try:
        model = ModelFactory.create_embedding_model(
            provider="openai",
            model_name="embeddinggemma",
            api_key="ollama",
            base_url="http://localhost:11434/v1",
        )

        texts = [
            "Hello from Ollama!",
            "This is a test embedding.",
            "Embedding models are useful for semantic search.",
        ]

        request = EmbeddingRequest(input=texts)

        print("📤 Sending embedding request to Ollama...")
        print("   Model: nomic-embed-text")
        print(f"   Texts to embed: {len(texts)}")
        for i, text in enumerate(texts, 1):
            print(f"     {i}. {text}")

        response = await model.embed(request)

        print("✅ Embedding successful")
        print(f"   Number of embeddings: {len(response.embeddings)}")
        print(f"   Embedding dimension: {len(response.embeddings[0])}")
        print(f"   First embedding (first 5 values): {response.embeddings[0][:5]}")

    except Exception as e:
        print(f"❌ Embedding failed: {e}")
        raise


@pytest.mark.asyncio
async def test_multimodal_vision() -> None:
    """Test multimodal vision capabilities with Ollama."""
    print("\n" + "=" * 60)
    print("TEST 9: Multimodal Vision (Image Understanding)")
    print("=" * 60)

    credentials = {
        "api_key": "ollama",
        "base_url": "http://localhost:11434/v1",
    }

    try:
        model = ModelFactory.create_model(
            provider="ollama",
            model_name="qwen3-vl:235b-cloud",
            credentials=credentials,
        )

        # Check if model supports vision
        print("📋 Model capabilities:")
        print(f"   Vision support: {model.supports_capability(ModelCapability.VISION)}")
        print(
            f"   Multimodal input: {model.supports_capability(ModelCapability.MULTIMODAL_INPUT)}"
        )

        # Create a message with image URL (using a public test image)
        image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/3/3a/Cat03.jpg/1200px-Cat03.jpg"

        content_parts = [
            ContentPart(
                content_type=ContentType.TEXT,
                content="What do you see in this image? Describe it briefly.",
            ),
            ContentPart(
                content_type=ContentType.IMAGE_URL,
                content=image_url,
                mime_type="image/jpeg",
            ),
        ]

        message = Message(role="user", content=content_parts)

        request = GenerateRequest(
            messages=[message],
            temperature=0.7,
            max_tokens=200,
        )

        print("\n📤 Sending multimodal request to Ollama...")
        print("   Model: gpt-oss:20b-cloud")
        print(f"   Message parts: {len(content_parts)}")
        print("     - Text: What do you see in this image?")
        print(f"     - Image: {image_url}")

        response = await model.generate(request)

        print("\n✅ Vision analysis successful")
        print(f"   Response: {response.content}")

    except Exception as e:
        print(f"❌ Vision test failed: {e}")
        raise


@pytest.mark.asyncio
async def test_multimodal_with_base64() -> None:
    """Test multimodal with base64 encoded image."""
    print("\n" + "=" * 60)
    print("TEST 10: Multimodal with Base64 Image")
    print("=" * 60)

    credentials = {
        "api_key": "ollama",
        "base_url": "http://localhost:11434/v1",
    }

    try:
        model = ModelFactory.create_model(
            provider="ollama",
            model_name="qwen3-vl:235b-cloud",
            credentials=credentials,
        )

        # For this test, we'll use a simple base64 encoded 1x1 pixel image
        # In real usage, you would encode an actual image
        base64_image = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="

        content_parts = [
            ContentPart(
                content_type=ContentType.TEXT,
                content="Analyze this image and tell me what you see.",
            ),
            ContentPart(
                content_type=ContentType.IMAGE_BASE64,
                content=base64_image,
                mime_type="image/png",
            ),
        ]

        message = Message(role="user", content=content_parts)

        request = GenerateRequest(
            messages=[message],
            temperature=0.5,
            max_tokens=100,
        )

        print("📤 Sending base64 image request to Ollama...")
        print("   Model: gpt-oss:20b-cloud")
        print(f"   Message parts: {len(content_parts)}")
        print("     - Text: Analyze this image")
        print(f"     - Image (base64): {len(base64_image)} characters")

        response = await model.generate(request)

        print("✅ Base64 image analysis successful")
        print(f"   Response: {response.content}")

    except Exception as e:
        print(f"❌ Base64 image test failed: {e}")
        raise


async def run_async_tests() -> None:
    """Run async tests."""
    try:
        # await test_text_generation()
        # await test_streaming_generation()
        # await test_embedding_generation()
        await test_multimodal_vision()
        await test_multimodal_with_base64()
    except Exception as e:
        print(f"\n❌ Async tests failed: {e}")
        raise


def main() -> None:
    """Run all tests."""
    print("\n")
    print("╔" + "=" * 58 + "╗")
    print("║" + " " * 58 + "║")
    print("║" + "  OLLAMA LLM PROVIDER INTEGRATION TESTS".center(58) + "║")
    print("║" + " " * 58 + "║")
    print("╚" + "=" * 58 + "╝")

    try:
        # Sync tests
        # test_ollama_model_creation()
        # test_model_configuration()
        # test_embedding_model_creation()
        # test_factory_provider_support()
        # test_openai_model_direct_creation()

        # Async tests (actual API calls)
        print("\n" + "=" * 60)
        print("Running async tests (actual API calls)...")
        print("=" * 60)
        asyncio.run(run_async_tests())

        print("\n" + "=" * 60)
        print("✅ ALL TESTS PASSED!")
        print("=" * 60 + "\n")

    except Exception as e:
        print("\n" + "=" * 60)
        print(f"❌ TESTS FAILED: {e}")
        print("=" * 60 + "\n")
        raise


if __name__ == "__main__":
    main()
