"""
Unit tests for context_models.py

Tests the core data models and their validation logic.
Updated to match v2.0 API.
"""

import pytest
from datetime import datetime
from typing import List
import uuid

from context_engineering.context_models import (
    ContextElement, ContextWindow, ContextSession, PromptTemplate,
    ContextType, PromptTemplateType, ContextQuality, OptimizationStatus,
    ContextAnalysis, OptimizationTask, MultimodalContext, RAGContext
)


class TestContextElement:
    """Test ContextElement data model."""

    def test_context_element_creation(self):
        """Test basic ContextElement creation."""
        element = ContextElement(
            content="Test content",
            type=ContextType.SYSTEM,
            priority=8
        )

        assert element.content == "Test content"
        assert element.type == ContextType.SYSTEM
        assert element.priority == 8
        assert isinstance(element.id, str)
        assert isinstance(element.created_at, datetime)

    def test_context_element_defaults(self):
        """Test ContextElement default values."""
        element = ContextElement()

        assert element.content == ""
        assert element.type == ContextType.USER
        assert element.priority == 5
        assert element.metadata == {}
        assert element.tags == []

    def test_token_count_property(self):
        """Test token count property."""
        element = ContextElement(content="This is a test message with some words")

        # Token count is based on word count * 1.3
        word_count = len(element.content.split())
        expected_tokens = word_count * 1.3
        assert element.token_count == expected_tokens

    def test_context_element_priority_range(self):
        """Test ContextElement priority values."""
        # Test valid priorities
        element_low = ContextElement(priority=1)
        element_high = ContextElement(priority=10)

        assert element_low.priority == 1
        assert element_high.priority == 10

    def test_context_element_serialization(self):
        """Test ContextElement can be serialized via to_dict()."""
        element = ContextElement(
            content="Test",
            type=ContextType.ASSISTANT,
            metadata={"test": True}
        )

        element_dict = element.to_dict()

        assert element_dict["content"] == "Test"
        assert element_dict["type"] == "assistant"
        assert element_dict["metadata"] == {"test": True}


class TestContextWindow:
    """Test ContextWindow data model."""

    def test_context_window_creation(self):
        """Test basic ContextWindow creation."""
        window = ContextWindow(max_tokens=4096)

        assert window.max_tokens == 4096
        assert window.elements == []
        assert window.current_tokens == 0
        assert window.reserved_tokens == 512  # default

    def test_add_element(self):
        """Test adding elements to context window."""
        window = ContextWindow(max_tokens=1000)
        element = ContextElement(content="Test content", priority=8)

        result = window.add_element(element)

        assert result == True
        assert len(window.elements) == 1
        assert window.elements[0] == element
        assert window.current_tokens > 0

    def test_remove_element(self):
        """Test removing elements from context window."""
        window = ContextWindow(max_tokens=1000)
        element = ContextElement(content="Test content")

        window.add_element(element)
        assert len(window.elements) == 1

        result = window.remove_element(element.id)
        assert result == True
        assert len(window.elements) == 0
        assert window.current_tokens == 0

    def test_utilization_ratio(self):
        """Test token utilization ratio calculation."""
        window = ContextWindow(max_tokens=1000)
        element = ContextElement(content="word " * 100)  # 100 words

        window.add_element(element)

        expected_ratio = window.current_tokens / window.max_tokens
        assert window.utilization_ratio == expected_ratio
        assert 0 <= window.utilization_ratio <= 1

    def test_available_tokens(self):
        """Test available tokens calculation (includes reserved_tokens)."""
        window = ContextWindow(max_tokens=1000, reserved_tokens=100)
        element = ContextElement(content="word " * 10)  # Some words

        window.add_element(element)

        # available_tokens = max_tokens - current_tokens - reserved_tokens
        expected = window.max_tokens - window.current_tokens - window.reserved_tokens
        assert window.available_tokens == expected

    def test_token_overflow_prevention(self):
        """Test that token overflow is prevented by add_element."""
        window = ContextWindow(max_tokens=100, reserved_tokens=50)
        # This creates an element with ~130 tokens (100 words * 1.3)
        large_element = ContextElement(content="word " * 100)

        # Should return False and not add the element
        result = window.add_element(large_element)
        assert result == False
        assert len(window.elements) == 0

    def test_optimize_for_tokens(self):
        """Test optimize_for_tokens removes low priority elements."""
        window = ContextWindow(max_tokens=100, reserved_tokens=10)

        low_priority = ContextElement(content="Low priority content", priority=2)
        high_priority = ContextElement(content="High priority content", priority=9)

        # Force add elements by directly appending
        window.elements.append(low_priority)
        window.elements.append(high_priority)

        result = window.optimize_for_tokens()

        # Should have optimization info
        assert "removed_elements" in result
        assert "tokens_saved" in result


class TestPromptTemplate:
    """Test PromptTemplate data model."""

    def test_prompt_template_creation(self):
        """Test basic PromptTemplate creation."""
        template = PromptTemplate(
            name="Test Template",
            description="A test template",
            template="Hello {name}, you are a {role}",
            variables=["name", "role"]
        )

        assert template.name == "Test Template"
        assert template.description == "A test template"
        assert "name" in template.variables
        assert "role" in template.variables
        assert template.type == PromptTemplateType.COMPLETION

    def test_template_variable_extraction(self):
        """Test variable extraction from template using extract_variables()."""
        template = PromptTemplate(
            name="Auto Extract",
            template="You are {role} with {years} years of experience in {field}."
        )

        # Use the extract_variables method
        extracted = template.extract_variables()
        expected_vars = {"role", "years", "field"}
        assert set(extracted) == expected_vars

    def test_template_render(self):
        """Test template rendering with variables."""
        template = PromptTemplate(
            name="Render Test",
            template="Hello {name}, you are a {role}",
            variables=["name", "role"]
        )

        rendered = template.render({"name": "Alice", "role": "developer"})
        assert rendered == "Hello Alice, you are a developer"

    def test_template_usage_count(self):
        """Test usage count default value."""
        template = PromptTemplate(name="Usage Test", template="Test")

        assert template.usage_count == 0

        # Manually increment (no method exists)
        template.usage_count += 1
        assert template.usage_count == 1


class TestContextSession:
    """Test ContextSession data model."""

    def test_context_session_creation(self):
        """Test basic ContextSession creation."""
        session = ContextSession(
            name="Test Session",
            description="A test session"
        )

        assert session.name == "Test Session"
        assert session.description == "A test session"
        assert session.windows == []
        assert isinstance(session.id, str)

    def test_create_window(self):
        """Test creating context windows in session."""
        session = ContextSession(name="Test")

        window = session.create_window(max_tokens=2000)

        assert len(session.windows) == 1
        assert window.max_tokens == 2000
        assert session.active_window_id == window.id

    def test_get_active_window(self):
        """Test getting active window."""
        session = ContextSession(name="Test")

        # No active window initially
        assert session.get_active_window() is None

        # Create a window
        window = session.create_window()

        # Now should return the window
        active = session.get_active_window()
        assert active is not None
        assert active.id == window.id


class TestContextAnalysis:
    """Test ContextAnalysis data model."""

    def test_context_analysis_creation(self):
        """Test ContextAnalysis creation."""
        analysis = ContextAnalysis(
            context_id="test-context-123",
            analysis_type="comprehensive"
        )

        assert analysis.context_id == "test-context-123"
        assert analysis.analysis_type == "comprehensive"
        assert analysis.quality_score == 0.0
        assert analysis.metrics == {}
        assert analysis.insights == []

    def test_analysis_with_data(self):
        """Test analysis with full data."""
        analysis = ContextAnalysis(
            context_id="test",
            analysis_type="quality",
            quality_score=85.5,
            metrics={"tokens": 100, "efficiency": 0.8},
            insights=["Good structure", "Clear content"],
            issues=["Minor redundancy"],
            recommendations=["Consider shortening"]
        )

        assert analysis.quality_score == 85.5
        assert len(analysis.insights) == 2
        assert len(analysis.issues) == 1
        assert len(analysis.recommendations) == 1

    def test_analysis_to_dict(self):
        """Test analysis serialization."""
        analysis = ContextAnalysis(
            context_id="test",
            analysis_type="basic",
            quality_score=75.0
        )

        result = analysis.to_dict()

        assert result["context_id"] == "test"
        assert result["analysis_type"] == "basic"
        assert result["quality_score"] == 75.0


class TestOptimizationTask:
    """Test OptimizationTask data model."""

    def test_optimization_task_creation(self):
        """Test OptimizationTask creation."""
        task = OptimizationTask(
            context_id="test-context",
            optimization_type="token_reduction"
        )

        assert task.context_id == "test-context"
        assert task.optimization_type == "token_reduction"
        assert task.status == OptimizationStatus.PENDING
        assert task.progress == 0.0

    def test_task_status_values(self):
        """Test task status enum values."""
        task = OptimizationTask(context_id="test")

        assert task.status == OptimizationStatus.PENDING

        # Change status manually
        task.status = OptimizationStatus.IN_PROGRESS
        assert task.status == OptimizationStatus.IN_PROGRESS

        task.status = OptimizationStatus.COMPLETED
        assert task.status == OptimizationStatus.COMPLETED

    def test_task_failure_status(self):
        """Test task failure status."""
        task = OptimizationTask(context_id="test")

        task.status = OptimizationStatus.FAILED
        task.error_message = "Test error message"

        assert task.status == OptimizationStatus.FAILED
        assert task.error_message == "Test error message"


class TestMultimodalContext:
    """Test MultimodalContext data model."""

    def test_multimodal_context_creation(self):
        """Test MultimodalContext creation."""
        context = MultimodalContext(
            text_content="Test text",
            image_urls=["http://example.com/image.png"]
        )

        assert context.text_content == "Test text"
        assert len(context.image_urls) == 1
        assert context.processing_status == "pending"

    def test_total_token_estimate(self):
        """Test token estimation for multimodal content."""
        context = MultimodalContext(
            text_content="This is some text content",
            image_urls=["img1.png", "img2.png"]
        )

        # Should include text tokens + image tokens (1000 per image)
        assert context.total_token_estimate > 2000


class TestRAGContext:
    """Test RAGContext data model."""

    def test_rag_context_creation(self):
        """Test RAGContext creation."""
        context = RAGContext(query="What is machine learning?")

        assert context.query == "What is machine learning?"
        assert context.retrieved_documents == []
        assert context.similarity_scores == []

    def test_add_retrieved_document(self):
        """Test adding retrieved documents."""
        context = RAGContext(query="test")

        doc = {"content": "Test document content", "source": "test.pdf"}
        context.add_retrieved_document(doc, 0.95)

        assert len(context.retrieved_documents) == 1
        assert len(context.similarity_scores) == 1
        assert context.similarity_scores[0] == 0.95

    def test_synthesize_context(self):
        """Test context synthesis from retrieved documents."""
        context = RAGContext(query="test")

        context.add_retrieved_document({"content": "Document A content"}, 0.9)
        context.add_retrieved_document({"content": "Document B content"}, 0.8)

        result = context.synthesize_context(max_tokens=1000)

        assert len(result) > 0
        assert context.context_synthesis == result


# Performance and edge case tests
class TestModelPerformance:
    """Test model performance and edge cases."""

    def test_large_context_window(self):
        """Test handling of large context windows."""
        window = ContextWindow(max_tokens=100000)  # Large window

        # Add many elements
        for i in range(100):
            element = ContextElement(content=f"Element {i}")
            window.add_element(element)

        assert len(window.elements) == 100
        assert window.current_tokens > 0

    def test_unicode_content_handling(self):
        """Test handling of Unicode content."""
        element = ContextElement(
            content="Unicode test: 日本語 العربية Ελληνικά"
        )

        assert "日本語" in element.content
        assert element.token_count > 0

    def test_empty_content(self):
        """Test handling of empty content."""
        element = ContextElement(content="")
        assert element.content == ""
        assert element.token_count == 0

    def test_token_calculation_speed(self):
        """Test token calculation is fast."""
        import time

        content = "This is a test content. " * 1000  # Large content

        start = time.time()
        element = ContextElement(content=content)
        tokens = element.token_count
        duration = time.time() - start

        assert tokens > 0
        assert duration < 0.1  # Should complete in less than 100ms