#!/usr/bin/env python3
"""
Demo script to test Ollama structured output with models that don't support format
"""

from pydantic import BaseModel
import json

# Import the updated Ollama implementation
from dazllm.llm_ollama import LlmOllama


class SimpleData(BaseModel):
    """Simple test data structure"""
    name: str
    value: int
    active: bool


def test_format_support_detection():
    """Test that we correctly detect format support"""
    print("Testing format support detection...")

    # Create mock instances to test detection
    test_llm = LlmOllama.__new__(LlmOllama)

    # Test model without format support
    test_llm.model = "gpt-oss:20b"
    supports_format = test_llm._check_format_support()
    print(f"  gpt-oss:20b supports format: {supports_format}")
    assert not supports_format, "gpt-oss:20b should NOT support format"

    # Test model with format support
    test_llm.model = "mistral-small"
    supports_format = test_llm._check_format_support()
    print(f"  mistral-small supports format: {supports_format}")
    assert supports_format, "mistral-small SHOULD support format"

    print("✓ Format support detection working correctly\n")


def test_json_parsing_fallbacks():
    """Test various JSON parsing strategies"""
    print("Testing JSON parsing fallbacks...")

    test_llm = LlmOllama.__new__(LlmOllama)

    test_cases = [
        ('```json\n{"name": "test", "value": 42, "active": true}\n```',
         {"name": "test", "value": 42, "active": True}),

        ('{"name": "plain", "value": 100, "active": false}',
         {"name": "plain", "value": 100, "active": False}),

        ('Here is the JSON: {"name": "prefixed", "value": 5, "active": true}',
         {"name": "prefixed", "value": 5, "active": True}),
    ]

    for input_text, expected in test_cases:
        try:
            result = test_llm._parse_json_with_fallbacks(input_text)
            assert result == expected, f"Expected {expected}, got {result}"
            print(f"  ✓ Parsed: {input_text[:30]}...")
        except Exception as e:
            print(f"  ✗ Failed to parse: {input_text[:30]}... - {e}")

    print("✓ JSON parsing fallbacks working correctly\n")


def test_schema_injection():
    """Test that schema is properly injected for non-format models"""
    print("Testing schema injection for non-format models...")

    # Create a simple schema
    schema_json = SimpleData.model_json_schema()
    schema_str = json.dumps(schema_json, indent=2)

    print(f"  Schema to inject:\n{schema_str[:200]}...")

    # Simulate what would happen with a model that doesn't support format
    system_prompt = (
        "You must respond with valid JSON that matches the following schema exactly. "
        "Output ONLY the JSON, wrapped in ```json code blocks. "
        "Do not include any other text, explanations, or the schema itself.\n\n"
        f"Required JSON Schema:\n{schema_str}"
    )

    print(f"\n  System prompt preview:\n{system_prompt[:200]}...")
    print("✓ Schema injection prepared correctly\n")


def main():
    """Run all demonstration tests"""
    print("=" * 60)
    print("Ollama Structured Output Demo")
    print("Testing support for models without 'format' parameter")
    print("=" * 60)
    print()

    try:
        test_format_support_detection()
        test_json_parsing_fallbacks()
        test_schema_injection()

        print("=" * 60)
        print("All tests passed! ✓")
        print("=" * 60)
        print("\nKey features implemented:")
        print("1. Auto-detection of models that don't support 'format' parameter")
        print("2. Fallback to manual schema injection in prompts")
        print("3. Multiple JSON parsing strategies with fallbacks")
        print("4. Retry logic for invalid JSON or schema mismatches")
        print("\nModels configured without format support:")
        for model in LlmOllama.MODELS_WITHOUT_FORMAT_SUPPORT:
            print(f"  - {model}")

    except AssertionError as e:
        print(f"\n✗ Test failed: {e}")
        return 1
    except Exception as e:
        print(f"\n✗ Unexpected error: {e}")
        return 1

    return 0


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