#!/usr/bin/env python3
"""Comprehensive Storage Feature Tests for Truthound.

This script tests all storage features manually to verify functionality.
"""

import asyncio
import json
import os
import sys
import tempfile
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any

# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))


def print_header(title: str) -> None:
    """Print a section header."""
    print(f"\n{'='*60}")
    print(f"  {title}")
    print(f"{'='*60}")


def print_subheader(title: str) -> None:
    """Print a subsection header."""
    print(f"\n--- {title} ---")


def print_result(name: str, success: bool, details: str = "") -> None:
    """Print a test result."""
    status = "PASS" if success else "FAIL"
    icon = "✓" if success else "✗"
    print(f"  [{status}] {icon} {name}")
    if details:
        print(f"        {details}")


class StorageFeatureTests:
    """Comprehensive storage feature tests."""

    def __init__(self):
        self.temp_dir = tempfile.mkdtemp(prefix="truthound_storage_test_")
        self.passed = 0
        self.failed = 0
        self.errors: list[str] = []

    def cleanup(self):
        """Clean up temporary files."""
        import shutil
        try:
            shutil.rmtree(self.temp_dir)
        except Exception:
            pass

    def run_test(self, name: str, test_func) -> bool:
        """Run a test and track results."""
        try:
            result = test_func()
            if result:
                self.passed += 1
                print_result(name, True)
            else:
                self.failed += 1
                print_result(name, False)
            return result
        except Exception as e:
            self.failed += 1
            error_msg = f"{type(e).__name__}: {e}"
            self.errors.append(f"{name}: {error_msg}")
            print_result(name, False, error_msg)
            return False

    # ========================================
    # 1. Factory and Backend Tests
    # ========================================

    def test_factory_list_backends(self) -> bool:
        """Test listing available backends."""
        from truthound.stores.factory import list_available_backends, is_backend_available

        backends = list_available_backends()
        # At minimum, filesystem and memory should be available
        assert "filesystem" in backends, "filesystem backend not available"
        assert "memory" in backends, "memory backend not available"

        assert is_backend_available("filesystem"), "filesystem should be available"
        assert is_backend_available("memory"), "memory should be available"
        assert not is_backend_available("nonexistent"), "nonexistent should not be available"

        return True

    def test_filesystem_store_basic(self) -> bool:
        """Test basic filesystem store operations."""
        from truthound.stores import get_store, ValidationResult, ResultStatus, ResultStatistics

        base_path = os.path.join(self.temp_dir, "fs_basic")
        store = get_store("filesystem", base_path=base_path)

        # Create a validation result
        result = ValidationResult(
            run_id="test_run_001",
            run_time=datetime.now(timezone.utc),
            data_asset="test_data.csv",
            status=ResultStatus.SUCCESS,  # Use SUCCESS instead of PASSED
            results=[],
            statistics=ResultStatistics(total_validators=10, passed_validators=10, failed_validators=0),
            tags={"env": "test"},
            metadata={"version": "1.0"},
        )

        # Save
        run_id = store.save(result)
        assert run_id, "Should return run_id"

        # Check exists
        assert store.exists(run_id), "Result should exist"

        # Get
        retrieved = store.get(run_id)
        assert retrieved is not None, "Should retrieve result"
        assert retrieved.data_asset == "test_data.csv", "Data asset should match"
        assert retrieved.status == ResultStatus.SUCCESS, "Status should match"

        # List
        run_ids = store.list_ids()
        assert len(run_ids) >= 1, "Should have at least one run"

        # Delete
        deleted = store.delete(run_id)
        assert deleted, "Should delete successfully"
        assert not store.exists(run_id), "Should not exist after delete"

        store.close()
        return True

    def test_memory_store_basic(self) -> bool:
        """Test basic memory store operations."""
        from truthound.stores import get_store, ValidationResult, ResultStatus, ResultStatistics

        store = get_store("memory")

        # Create and save
        result = ValidationResult(
            run_id="mem_test_001",
            run_time=datetime.now(timezone.utc),
            data_asset="memory_test.csv",
            status=ResultStatus.FAILURE,  # Use FAILURE instead of FAILED
            results=[],
            statistics=ResultStatistics(total_validators=5, passed_validators=3, failed_validators=2),
            tags={},
            metadata={},
        )

        run_id = store.save(result)
        assert run_id, "Should return run_id"

        # Verify
        retrieved = store.get(run_id)
        assert retrieved.status == ResultStatus.FAILURE, "Status should be FAILURE"

        # Clear all
        count = store.clear()
        assert count >= 1, "Should clear at least one item"

        store.close()
        return True

    def test_store_query(self) -> bool:
        """Test store query functionality."""
        from truthound.stores import get_store, ValidationResult, ResultStatus, ResultStatistics, StoreQuery

        store = get_store("memory")

        # Create multiple results
        for i in range(5):
            result = ValidationResult(
                run_id=f"query_test_{i:03d}",
                run_time=datetime.now(timezone.utc) - timedelta(days=i),
                data_asset="query_data.csv" if i < 3 else "other_data.csv",
                status=ResultStatus.SUCCESS if i % 2 == 0 else ResultStatus.FAILURE,
                results=[],
                statistics=ResultStatistics(total_validators=10, passed_validators=10 - i, failed_validators=i),
                tags={"batch": str(i // 2)},
                metadata={},
            )
            store.save(result)

        # Query by data_asset
        query = StoreQuery(data_asset="query_data.csv")
        results = store.query(query)
        assert len(results) == 3, f"Should find 3 results for query_data.csv, got {len(results)}"

        # Query with limit
        query = StoreQuery(limit=2)
        results = store.query(query)
        assert len(results) == 2, f"Should find 2 results with limit, got {len(results)}"

        # Count
        total = store.count()
        assert total == 5, f"Should have 5 total results, got {total}"

        store.close()
        return True

    # ========================================
    # 2. Versioning Tests
    # ========================================

    def test_versioning_incremental(self) -> bool:
        """Test incremental versioning strategy."""
        from truthound.stores.versioning import (
            IncrementalStrategy,
            VersioningConfig,
            VersionedStore,
        )
        from truthound.stores import get_store, ValidationResult, ResultStatus, ResultStatistics

        base_store = get_store("memory")
        config = VersioningConfig(strategy=IncrementalStrategy(), max_versions=5)
        versioned = VersionedStore(base_store, config)

        # Save multiple versions
        versions = []
        for i in range(3):
            result = ValidationResult(
                run_id=f"version_test_{i}",
                run_time=datetime.now(timezone.utc),
                data_asset="versioned_data.csv",
                status=ResultStatus.SUCCESS,
                results=[],
                statistics=ResultStatistics(total_validators=i+1),
                tags={},
                metadata={"version_iteration": i},
            )
            run_id = versioned.save(result)
            versions.append(run_id)

        # Check version history
        history = versioned.get_version_history("versioned_data.csv")
        assert len(history) == 3, f"Should have 3 versions, got {len(history)}"

        versioned.close()
        return True

    def test_versioning_semantic(self) -> bool:
        """Test semantic versioning strategy."""
        from truthound.stores.versioning import (
            SemanticStrategy,
            VersioningConfig,
            VersionedStore,
        )
        from truthound.stores import get_store, ValidationResult, ResultStatus, ResultStatistics

        base_store = get_store("memory")
        strategy = SemanticStrategy()
        config = VersioningConfig(strategy=strategy, max_versions=10)
        versioned = VersionedStore(base_store, config)

        # Save versions
        for i in range(3):
            result = ValidationResult(
                run_id=f"semantic_test_{i}",
                run_time=datetime.now(timezone.utc),
                data_asset="semantic_data.csv",
                status=ResultStatus.SUCCESS,
                results=[],
                statistics=ResultStatistics(),
                tags={},
                metadata={},
            )
            versioned.save(result)

        # Check version format
        history = versioned.get_version_history("semantic_data.csv")
        assert len(history) >= 1, "Should have at least one version"

        versioned.close()
        return True

    def test_versioning_timestamp(self) -> bool:
        """Test timestamp versioning strategy."""
        from truthound.stores.versioning import (
            TimestampStrategy,
            VersioningConfig,
            VersionedStore,
        )
        from truthound.stores import get_store, ValidationResult, ResultStatus, ResultStatistics

        base_store = get_store("memory")
        strategy = TimestampStrategy()
        config = VersioningConfig(strategy=strategy)
        versioned = VersionedStore(base_store, config)

        result = ValidationResult(
            run_id="timestamp_test_001",
            run_time=datetime.now(timezone.utc),
            data_asset="timestamp_data.csv",
            status=ResultStatus.SUCCESS,
            results=[],
            statistics=ResultStatistics(),
            tags={},
            metadata={},
        )
        versioned.save(result)

        history = versioned.get_version_history("timestamp_data.csv")
        assert len(history) == 1, "Should have one version"

        versioned.close()
        return True

    def test_versioning_gitlike(self) -> bool:
        """Test git-like versioning strategy."""
        from truthound.stores.versioning import (
            GitLikeStrategy,
            VersioningConfig,
            VersionedStore,
        )
        from truthound.stores import get_store, ValidationResult, ResultStatus, ResultStatistics

        base_store = get_store("memory")
        strategy = GitLikeStrategy()
        config = VersioningConfig(strategy=strategy)
        versioned = VersionedStore(base_store, config)

        result = ValidationResult(
            run_id="gitlike_test_001",
            run_time=datetime.now(timezone.utc),
            data_asset="gitlike_data.csv",
            status=ResultStatus.SUCCESS,
            results=[],
            statistics=ResultStatistics(total_validators=5),
            tags={},
            metadata={},
        )
        versioned.save(result)

        history = versioned.get_version_history("gitlike_data.csv")
        assert len(history) == 1, "Should have one version"

        versioned.close()
        return True

    # ========================================
    # 3. Retention Policy Tests
    # ========================================

    def test_retention_time_based(self) -> bool:
        """Test time-based retention policy."""
        from truthound.stores.retention import TimeBasedPolicy
        from truthound.stores import ValidationResult, ResultStatus, ResultStatistics

        policy = TimeBasedPolicy(max_age_days=30)

        # Create old and new results
        old_result = ValidationResult(
            run_id="old_result",
            run_time=datetime.now(timezone.utc) - timedelta(days=60),  # 60 days old
            data_asset="retention_test.csv",
            status=ResultStatus.SUCCESS,
            results=[],
            statistics=ResultStatistics(),
            tags={},
            metadata={},
        )
        new_result = ValidationResult(
            run_id="new_result",
            run_time=datetime.now(timezone.utc),  # Current
            data_asset="retention_test.csv",
            status=ResultStatus.SUCCESS,
            results=[],
            statistics=ResultStatistics(),
            tags={},
            metadata={},
        )

        # Check which should be retained
        should_retain_old = policy.should_retain(old_result)
        should_retain_new = policy.should_retain(new_result)

        assert not should_retain_old, "Old result should not be retained"
        assert should_retain_new, "New result should be retained"

        return True

    def test_retention_count_based(self) -> bool:
        """Test count-based retention policy."""
        from truthound.stores.retention import CountBasedPolicy
        from truthound.stores import ValidationResult, ResultStatus, ResultStatistics

        policy = CountBasedPolicy(max_count=3)

        results = []
        for i in range(5):
            result = ValidationResult(
                run_id=f"count_test_{i}",
                run_time=datetime.now(timezone.utc) - timedelta(hours=i),
                data_asset="count_test.csv",
                status=ResultStatus.SUCCESS,
                results=[],
                statistics=ResultStatistics(),
                tags={},
                metadata={},
            )
            results.append(result)

        # CountBasedPolicy stores max_count internally
        assert policy.max_count == 3, "Max count should be 3"

        return True

    def test_retention_size_based(self) -> bool:
        """Test size-based retention policy."""
        from truthound.stores.retention import SizeBasedPolicy

        policy = SizeBasedPolicy(max_size_mb=100)  # 100 MB

        assert policy.max_size_mb == 100, "Max size should be 100 MB"

        return True

    def test_retention_composite(self) -> bool:
        """Test composite retention policy."""
        from truthound.stores.retention import (
            TimeBasedPolicy,
            CountBasedPolicy,
            CompositePolicy,
        )

        time_policy = TimeBasedPolicy(max_age_days=30)
        count_policy = CountBasedPolicy(max_count=10)
        composite = CompositePolicy(policies=[time_policy, count_policy])

        assert len(composite.policies) == 2, "Should have 2 policies"

        return True

    # ========================================
    # 4. Tiered Storage Tests
    # ========================================

    def test_tiering_basic(self) -> bool:
        """Test basic tiered storage functionality."""
        from truthound.stores.tiering import (
            StorageTier,
            TierType,
            TierInfo,
            InMemoryTierMetadataStore,
        )

        # Create tiers
        hot = StorageTier(name="hot", tier_type=TierType.HOT, priority=1)
        warm = StorageTier(name="warm", tier_type=TierType.WARM, priority=2)
        cold = StorageTier(name="cold", tier_type=TierType.COLD, priority=3)

        assert hot.priority < warm.priority, "Hot should have lower priority than warm"
        assert warm.priority < cold.priority, "Warm should have lower priority than cold"

        # Test tier metadata store
        metadata_store = InMemoryTierMetadataStore()
        info = TierInfo(
            item_id="test_item",
            tier_name="hot",
            created_at=datetime.now(timezone.utc),
            access_count=0,
        )
        metadata_store.save(info)

        retrieved = metadata_store.get("test_item")
        assert retrieved is not None, "Should retrieve tier info"
        assert retrieved.tier_name == "hot", "Should be in hot tier"

        return True

    def test_tiering_age_based_policy(self) -> bool:
        """Test age-based tier migration policy."""
        from truthound.stores.tiering import AgeBasedTierPolicy, TierInfo

        # Create policy with age thresholds
        policy = AgeBasedTierPolicy(
            days_to_warm=7,
            days_to_cold=30,
            days_to_archive=90,
        )

        # Test hot item (new)
        hot_info = TierInfo(
            item_id="hot_item",
            tier_name="hot",
            created_at=datetime.now(timezone.utc),
            access_count=0,
        )

        # Test warm candidate (8 days old)
        warm_info = TierInfo(
            item_id="warm_item",
            tier_name="hot",
            created_at=datetime.now(timezone.utc) - timedelta(days=8),
            access_count=0,
        )

        hot_target = policy.should_migrate(hot_info)
        warm_target = policy.should_migrate(warm_info)

        assert hot_target is None, "Hot item should not migrate"
        assert warm_target is not None, "Warm candidate should migrate"

        return True

    def test_tiering_access_based_policy(self) -> bool:
        """Test access-based tier migration policy."""
        from truthound.stores.tiering import AccessBasedTierPolicy, TierInfo

        policy = AccessBasedTierPolicy(
            access_threshold_promote=10,  # Promote if accessed 10+ times
            access_threshold_demote=1,    # Demote if accessed <= 1 time
        )

        # Frequently accessed item in warm tier
        frequently_accessed = TierInfo(
            item_id="freq_item",
            tier_name="warm",
            created_at=datetime.now(timezone.utc) - timedelta(days=5),
            access_count=15,
            last_accessed=datetime.now(timezone.utc),
        )

        # Rarely accessed item in hot tier
        rarely_accessed = TierInfo(
            item_id="rare_item",
            tier_name="hot",
            created_at=datetime.now(timezone.utc) - timedelta(days=10),
            access_count=1,
            last_accessed=datetime.now(timezone.utc) - timedelta(days=8),
        )

        freq_target = policy.should_migrate(frequently_accessed)
        rare_target = policy.should_migrate(rarely_accessed)

        # Verify the policy returns migration targets
        assert freq_target is not None or rarely_accessed.access_count <= policy.access_threshold_demote, "Should evaluate access patterns"

        return True

    # ========================================
    # 5. Caching Tests
    # ========================================

    def test_cache_inmemory(self) -> bool:
        """Test in-memory cache."""
        from truthound.stores.caching import InMemoryCache, CacheConfig

        config = CacheConfig(max_size=100, ttl_seconds=60.0)
        cache = InMemoryCache(config)

        # Set and get
        cache.set("key1", {"data": "value1"})
        value = cache.get("key1")
        assert value == {"data": "value1"}, f"Should retrieve value, got {value}"

        # Delete
        cache.delete("key1")
        assert cache.get("key1") is None, "Key should not exist after delete"

        # Clear
        cache.set("key2", "value2")
        cache.set("key3", "value3")
        cache.clear()
        assert cache.get("key2") is None, "Cache should be cleared"

        return True

    def test_cache_lru(self) -> bool:
        """Test LRU cache eviction."""
        from truthound.stores.caching import LRUCache

        cache = LRUCache(max_size=3)

        # Fill cache
        cache.set("a", 1)
        cache.set("b", 2)
        cache.set("c", 3)

        # Access 'a' to make it recently used
        _ = cache.get("a")

        # Add new item, should evict 'b' (least recently used)
        cache.set("d", 4)

        assert cache.get("a") is not None, "'a' should still exist (recently used)"
        assert cache.get("b") is None, "'b' should be evicted (LRU)"
        assert cache.get("c") is not None, "'c' should still exist"
        assert cache.get("d") is not None, "'d' should exist (new)"

        return True

    def test_cache_lfu(self) -> bool:
        """Test LFU cache eviction."""
        from truthound.stores.caching import LFUCache

        cache = LFUCache(max_size=3)

        # Fill cache with different access frequencies
        cache.set("a", 1)
        cache.set("b", 2)
        cache.set("c", 3)

        # Access 'a' multiple times
        for _ in range(5):
            cache.get("a")

        # Access 'b' a few times
        for _ in range(3):
            cache.get("b")

        # 'c' accessed only once (during set)

        # Add new item, should evict 'c' (least frequently used)
        cache.set("d", 4)

        assert cache.get("a") is not None, "'a' should still exist (most frequent)"
        assert cache.get("b") is not None, "'b' should still exist (moderate frequency)"

        return True

    def test_cache_ttl(self) -> bool:
        """Test TTL cache expiration."""
        from truthound.stores.caching import TTLCache

        cache = TTLCache(ttl_seconds=0.1, max_size=100)  # 100ms TTL

        cache.set("short_lived", "value")
        assert cache.get("short_lived") is not None, "Should exist immediately"

        # Wait for expiration
        time.sleep(0.15)

        assert cache.get("short_lived") is None, "Should expire after TTL"

        return True

    def test_cached_store(self) -> bool:
        """Test CachedStore wrapper."""
        from truthound.stores.caching import CachedStore, LRUCache, CacheMode
        from truthound.stores import get_store, ValidationResult, ResultStatus, ResultStatistics

        base_store = get_store("memory")
        cache = LRUCache(max_size=100)
        cached = CachedStore(base_store, cache, mode=CacheMode.WRITE_THROUGH)

        # Save through cache
        result = ValidationResult(
            run_id="cached_test_001",
            run_time=datetime.now(timezone.utc),
            data_asset="cached_data.csv",
            status=ResultStatus.SUCCESS,
            results=[],
            statistics=ResultStatistics(),
            tags={},
            metadata={},
        )
        run_id = cached.save(result)

        # First get should hit cache
        retrieved = cached.get(run_id)
        assert retrieved is not None, "Should retrieve from cache"

        cached.close()
        return True

    # ========================================
    # 6. Batch Write Tests
    # ========================================

    def test_batch_buffer(self) -> bool:
        """Test batch buffer functionality."""
        from truthound.stores.batching import BatchBuffer

        buffer = BatchBuffer(max_size=5)

        # Add items
        for i in range(3):
            added = buffer.add(f"item_{i}")
            assert added, f"Should add item {i}"

        assert len(buffer) == 3, f"Should have 3 items, got {len(buffer)}"

        # Fill buffer
        buffer.add("item_3")
        buffer.add("item_4")

        # Drain
        items = buffer.drain(2)
        assert len(items) == 2, f"Should drain 2 items, got {len(items)}"
        assert len(buffer) == 3, f"Should have 3 items remaining, got {len(buffer)}"

        # Drain all
        items = buffer.drain_all()
        assert len(items) == 3, f"Should drain all 3 items, got {len(items)}"
        assert len(buffer) == 0, "Buffer should be empty"

        return True

    def test_memory_aware_buffer(self) -> bool:
        """Test memory-aware buffer."""
        from truthound.stores.batching import MemoryAwareBuffer

        buffer = MemoryAwareBuffer(
            max_size=100,
            max_memory_mb=1,  # 1 MB
        )

        # Add items
        for i in range(10):
            buffer.add(f"item_{i}" * 10)  # Each item ~60 bytes

        stats = buffer.get_stats()
        assert stats["item_count"] > 0, "Should have items"
        assert stats["memory_used"] > 0, "Should track memory"

        return True

    def test_batch_writer(self) -> bool:
        """Test batch writer."""
        from truthound.stores.batching import BatchWriter, BatchConfig

        written_batches = []

        def write_func(batch: list) -> None:
            written_batches.append(batch)

        config = BatchConfig(batch_size=3, max_pending_batches=10)
        writer = BatchWriter(write_func, config)

        # Add items
        with writer:
            for i in range(7):
                writer.add(f"item_{i}")

        # Should have 3 batches: [0,1,2], [3,4,5], [6]
        assert len(written_batches) == 3, f"Should have 3 batches, got {len(written_batches)}"
        assert len(written_batches[0]) == 3, "First batch should have 3 items"
        assert len(written_batches[1]) == 3, "Second batch should have 3 items"
        assert len(written_batches[2]) == 1, "Third batch should have 1 item"

        return True

    def test_async_batch_writer(self) -> bool:
        """Test async batch writer."""
        from truthound.stores.batching import AsyncBatchWriter, BatchConfig

        written_batches = []

        async def async_write_func(batch: list) -> None:
            written_batches.append(batch)
            await asyncio.sleep(0.01)  # Simulate async operation

        async def run_test():
            config = BatchConfig(batch_size=2, parallelism=2)
            writer = AsyncBatchWriter(async_write_func, config)

            async with writer:
                for i in range(5):
                    await writer.add(f"async_item_{i}")

            return len(written_batches)

        batch_count = asyncio.run(run_test())
        assert batch_count == 3, f"Should have 3 batches, got {batch_count}"

        return True

    # ========================================
    # 7. Backpressure Tests
    # ========================================

    def test_backpressure_memory(self) -> bool:
        """Test memory-based backpressure."""
        from truthound.stores.backpressure import (
            MemoryBasedBackpressure,
            BackpressureConfig,
            PressureLevel,
        )

        config = BackpressureConfig(
            memory_threshold_percent=80.0,
        )

        # Mock memory provider - normal state
        backpressure = MemoryBasedBackpressure(
            config, memory_provider=lambda: 50.0
        )

        level = backpressure.check_pressure()
        assert level == PressureLevel.NONE, f"Should be NONE at 50%, got {level}"

        # High pressure state
        backpressure = MemoryBasedBackpressure(
            config, memory_provider=lambda: 85.0
        )
        level = backpressure.check_pressure()
        assert level in (PressureLevel.MEDIUM, PressureLevel.HIGH), f"Should be elevated at 85%, got {level}"

        return True

    def test_backpressure_queue_depth(self) -> bool:
        """Test queue-depth backpressure."""
        from truthound.stores.backpressure import (
            QueueDepthBackpressure,
            BackpressureConfig,
            PressureLevel,
        )

        config = BackpressureConfig(
            queue_depth_threshold=100,
        )
        backpressure = QueueDepthBackpressure(config)

        # Set queue depth
        backpressure.update_queue_depth(50)
        level = backpressure.check_pressure()
        assert level == PressureLevel.NONE, f"Should be NONE at 50 items, got {level}"

        backpressure.update_queue_depth(110)
        level = backpressure.check_pressure()
        assert level != PressureLevel.NONE, f"Should have pressure at 110 items, got {level}"

        return True

    def test_backpressure_token_bucket(self) -> bool:
        """Test token bucket rate limiter."""
        from truthound.stores.backpressure import (
            TokenBucketBackpressure,
            BackpressureConfig,
        )

        config = BackpressureConfig(
            tokens_per_second=10.0,  # 10 tokens per second
            bucket_size=5,           # 5 token burst
        )
        backpressure = TokenBucketBackpressure(config)

        # Should be able to acquire burst
        acquired_count = 0
        for i in range(5):
            if backpressure.try_acquire():
                acquired_count += 1

        assert acquired_count > 0, "Should acquire at least some tokens"

        return True

    def test_circuit_breaker(self) -> bool:
        """Test circuit breaker pattern."""
        from truthound.stores.backpressure import CircuitBreaker, CircuitBreakerConfig

        config = CircuitBreakerConfig(
            failure_threshold=3,
            success_threshold=2,
            timeout_seconds=0.1,  # 100ms for testing
        )
        breaker = CircuitBreaker(config)

        # Initially closed
        assert breaker.is_closed, "Should start closed"

        # Record failures using context manager
        for _ in range(3):
            try:
                with breaker:
                    raise Exception("Simulated failure")
            except Exception:
                pass

        # Should be open now
        assert breaker.is_open, "Should be open after 3 failures"

        # Wait for timeout
        time.sleep(0.15)

        # Should be half-open
        assert breaker.is_half_open, "Should be half-open after timeout"

        return True

    # ========================================
    # 8. Replication Tests
    # ========================================

    def test_replication_basic(self) -> bool:
        """Test basic replication setup."""
        from truthound.stores.replication import (
            ReplicationConfig,
            ReplicationMode,
            ReplicaTarget,
        )
        from truthound.stores import get_store

        primary = get_store("memory")
        replica1 = get_store("memory")

        # Create replica target
        target = ReplicaTarget(name="replica1", store=replica1)

        config = ReplicationConfig(
            mode=ReplicationMode.SYNC,
            targets=[target],
        )

        assert config.mode == ReplicationMode.SYNC, "Mode should be SYNC"
        assert len(config.targets) == 1, "Should have 1 replica target"

        primary.close()
        replica1.close()
        return True

    def test_replication_sync(self) -> bool:
        """Test synchronous replication."""
        from truthound.stores.replication import (
            ReplicationConfig,
            ReplicationMode,
            ReplicatedStore,
            ReplicaTarget,
        )
        from truthound.stores import get_store, ValidationResult, ResultStatus, ResultStatistics

        primary = get_store("memory")
        replica = get_store("memory")

        target = ReplicaTarget(name="replica1", store=replica)
        config = ReplicationConfig(
            mode=ReplicationMode.SYNC,
            targets=[target],
        )
        replicated = ReplicatedStore(primary, config)

        # Save to replicated store
        result = ValidationResult(
            run_id="replicated_test_001",
            run_time=datetime.now(timezone.utc),
            data_asset="replicated_data.csv",
            status=ResultStatus.SUCCESS,
            results=[],
            statistics=ResultStatistics(),
            tags={},
            metadata={},
        )
        run_id = replicated.save(result)

        # Verify exists in primary
        assert primary.exists(run_id), "Should exist in primary"

        # Verify exists in replica (sync mode)
        assert replica.exists(run_id), "Should exist in replica (sync)"

        replicated.close()
        return True

    # ========================================
    # 9. Migration Tests
    # ========================================

    def test_migration_registry(self) -> bool:
        """Test migration registry."""
        from truthound.stores.migration import MigrationRegistry, SchemaMigration

        registry = MigrationRegistry()

        def migrate_100_to_110(data: dict) -> dict:
            data["new_field"] = "default_value"
            return data

        def migrate_110_to_200(data: dict) -> dict:
            data["version"] = 2
            return data

        # Register migrations
        migration1 = SchemaMigration(
            from_version="1.0.0",
            to_version="1.1.0",
            migrate=migrate_100_to_110,
        )
        migration2 = SchemaMigration(
            from_version="1.1.0",
            to_version="2.0.0",
            migrate=migrate_110_to_200,
        )
        registry.register(migration1)
        registry.register(migration2)

        # Check migrations registered
        assert registry.has_migration("1.0.0", "1.1.0"), "Should have 1.0.0 -> 1.1.0"
        assert registry.has_migration("1.1.0", "2.0.0"), "Should have 1.1.0 -> 2.0.0"

        # Find migration path
        path = registry.find_path("1.0.0", "2.0.0")
        assert path is not None, "Should find migration path"
        assert len(path) == 2, f"Path should have 2 steps, got {len(path)}"

        return True

    def test_migration_execution(self) -> bool:
        """Test migration execution."""
        from truthound.stores.migration import MigrationRegistry, SchemaMigration, MigrationManager

        registry = MigrationRegistry()

        def upgrade_to_v2(data: dict) -> dict:
            data["upgraded"] = True
            data["schema_version"] = "2.0.0"
            return data

        migration = SchemaMigration(
            from_version="1.0.0",
            to_version="2.0.0",
            migrate=upgrade_to_v2,
        )
        registry.register(migration)

        manager = MigrationManager(registry)

        # Migrate data
        old_data = {"value": 123, "schema_version": "1.0.0"}
        new_data = manager.migrate(old_data, "1.0.0", "2.0.0")

        assert new_data["upgraded"] is True, "Should be upgraded"
        assert new_data["schema_version"] == "2.0.0", "Should have new version"
        assert new_data["value"] == 123, "Should preserve original data"

        return True

    # ========================================
    # 10. Encryption Tests
    # ========================================

    def test_encryption_basic(self) -> bool:
        """Test basic encryption."""
        from truthound.stores.encryption import (
            get_encryptor,
            generate_key,
            EncryptionAlgorithm,
        )

        key = generate_key(EncryptionAlgorithm.AES_256_GCM)
        encryptor = get_encryptor("aes-256-gcm")

        plaintext = b"This is sensitive validation data"
        encrypted = encryptor.encrypt(plaintext, key)

        assert encrypted != plaintext, "Encrypted should differ from plaintext"

        decrypted = encryptor.decrypt(encrypted, key)
        assert decrypted == plaintext, "Decrypted should match plaintext"

        return True

    def test_encryption_algorithms(self) -> bool:
        """Test multiple encryption algorithms."""
        from truthound.stores.encryption import (
            get_encryptor,
            generate_key,
            EncryptionAlgorithm,
        )

        # Test AES-256-GCM (always available)
        key = generate_key(EncryptionAlgorithm.AES_256_GCM)
        encryptor = get_encryptor("aes-256-gcm")

        plaintext = b"Test data for encryption verification"
        encrypted = encryptor.encrypt(plaintext, key)
        decrypted = encryptor.decrypt(encrypted, key)

        assert decrypted == plaintext, "AES encryption failed roundtrip"

        return True

    def test_encryption_pipeline(self) -> bool:
        """Test encryption pipeline with compression."""
        from truthound.stores.encryption import (
            create_secure_pipeline,
            generate_key,
            EncryptionAlgorithm,
        )

        key = generate_key(EncryptionAlgorithm.AES_256_GCM)
        pipeline = create_secure_pipeline(key, compression="gzip")

        original_data = b"This is test data " * 100  # Compressible data

        result = pipeline.process(original_data)
        recovered = pipeline.reverse(result.data)

        assert recovered == original_data, "Pipeline should roundtrip correctly"

        return True

    # ========================================
    # 11. Compression Tests
    # ========================================

    def test_compression_basic(self) -> bool:
        """Test basic compression."""
        from truthound.stores.compression import get_compressor

        compressor = get_compressor("gzip")

        original = b"This is test data that should compress well. " * 100
        compressed = compressor.compress(original)

        assert len(compressed) < len(original), "Compressed should be smaller"

        decompressed = compressor.decompress(compressed)
        assert decompressed == original, "Decompressed should match original"

        return True

    def test_compression_algorithms(self) -> bool:
        """Test gzip compression (always available)."""
        from truthound.stores.compression import get_compressor

        original = b"Compression test data " * 50

        compressor = get_compressor("gzip")
        compressed = compressor.compress(original)
        decompressed = compressor.decompress(compressed)
        assert decompressed == original, "gzip compression failed roundtrip"

        return True

    def test_compression_pipeline(self) -> bool:
        """Test compression pipeline."""
        from truthound.stores.compression import create_json_pipeline

        pipeline = create_json_pipeline()

        original = b'{"key": "value", "numbers": [1, 2, 3, 4, 5]}' * 100

        result = pipeline.process(original)
        recovered = pipeline.reverse(result.data)

        assert recovered == original, "Pipeline should roundtrip correctly"

        return True

    # ========================================
    # 12. Observability Tests
    # ========================================

    def test_audit_logging(self) -> bool:
        """Test audit logging functionality."""
        from truthound.stores.observability import (
            AuditEvent,
            AuditEventType,
            InMemoryAuditBackend,
            AuditLogger,
        )

        backend = InMemoryAuditBackend()
        logger = AuditLogger(backend)

        # Log an event
        logger.log_event(
            AuditEventType.CREATE,
            item_id="test_item_001",
            data_asset="audit_test.csv",
            metadata={"size": 1024},
        )

        # Query events
        events = backend.query()
        assert len(events) >= 1, "Should have at least 1 event"

        event = events[0]
        assert event.event_type == AuditEventType.CREATE, "Should be CREATE event"
        assert event.item_id == "test_item_001", "Should have correct item_id"

        return True

    def test_metrics_collection(self) -> bool:
        """Test metrics collection."""
        from truthound.stores.observability import (
            InMemoryMetricsBackend,
            StoreMetrics,
            MetricsConfig,
        )

        config = MetricsConfig(prefix="test_")
        backend = InMemoryMetricsBackend(config)
        metrics = StoreMetrics(backend, store_name="test_store", store_type="memory")

        # Record operations
        metrics.record_operation("save", latency=0.05, success=True)
        metrics.record_operation("get", latency=0.02, success=True)

        # Check metrics
        exported = backend.export()
        assert len(exported) > 0, "Should have metrics"

        return True

    def test_observable_store(self) -> bool:
        """Test observable store wrapper."""
        from truthound.stores.observability import (
            ObservableStore,
            ObservabilityConfig,
            AuditConfig,
            MetricsConfig,
        )
        from truthound.stores import get_store, ValidationResult, ResultStatus, ResultStatistics

        base_store = get_store("memory")
        config = ObservabilityConfig(
            audit=AuditConfig(enabled=True),
            metrics=MetricsConfig(enabled=True),
        )
        observable = ObservableStore(base_store, config)

        # Perform operations
        result = ValidationResult(
            run_id="observable_test_001",
            run_time=datetime.now(timezone.utc),
            data_asset="observable_data.csv",
            status=ResultStatus.SUCCESS,
            results=[],
            statistics=ResultStatistics(),
            tags={},
            metadata={},
        )
        run_id = observable.save(result)
        _ = observable.get(run_id)

        # Check audit events
        events = observable.get_audit_events()
        assert len(events) >= 2, f"Should have at least 2 events, got {len(events)}"

        observable.close()
        return True

    def test_tracing(self) -> bool:
        """Test distributed tracing."""
        from truthound.stores.observability import (
            SpanContext,
            NoopTracer,
        )

        tracer = NoopTracer()

        with tracer.trace("test_operation") as span:
            span.set_attribute("key", "value")
            span.add_event("processing_started")

            # Nested span
            with tracer.trace("sub_operation") as child:
                child.set_attribute("nested", True)

        # NoopTracer doesn't track spans but should work without error
        return True

    # ========================================
    # 13. Streaming Storage Tests
    # ========================================

    def test_streaming_filesystem(self) -> bool:
        """Test streaming filesystem store."""
        from truthound.stores.streaming import StreamingFileSystemStore
        from truthound.stores import ValidatorResult

        base_path = os.path.join(self.temp_dir, "streaming")
        store = StreamingFileSystemStore(base_path=base_path)

        # Create a session
        session = store.create_session("stream_run_001", "large_data.csv")

        # Write results incrementally
        with store.create_writer(session) as writer:
            for i in range(100):
                result = ValidatorResult(
                    validator_name=f"validator_{i}",
                    success=i % 2 == 0,
                    message=f"Result {i}",
                    details={"index": i},
                )
                writer.write_result(result)

        # Read back
        count = 0
        for result in store.iter_results("stream_run_001"):
            count += 1

        assert count == 100, f"Should read 100 results, got {count}"

        store.close()
        return True

    # ========================================
    # 14. Concurrency Tests
    # ========================================

    def test_concurrent_filesystem(self) -> bool:
        """Test concurrent filesystem store."""
        from truthound.stores.backends.concurrent_filesystem import (
            ConcurrentFileSystemStore,
            ConcurrencyConfig,
        )
        from truthound.stores import ValidationResult, ResultStatus, ResultStatistics

        base_path = os.path.join(self.temp_dir, "concurrent")
        config = ConcurrencyConfig(lock_strategy="file")
        store = ConcurrentFileSystemStore(base_path=base_path, concurrency=config)

        # Save result with locking
        result = ValidationResult(
            run_id="concurrent_test_001",
            run_time=datetime.now(timezone.utc),
            data_asset="concurrent_data.csv",
            status=ResultStatus.SUCCESS,
            results=[],
            statistics=ResultStatistics(),
            tags={},
            metadata={},
        )
        run_id = store.save(result)

        # Retrieve
        retrieved = store.get(run_id)
        assert retrieved is not None, "Should retrieve result"

        store.close()
        return True

    # ========================================
    # 15. Store Composition Tests
    # ========================================

    def test_store_composition(self) -> bool:
        """Test composing multiple store features."""
        from truthound.stores import get_store, ValidationResult, ResultStatus, ResultStatistics
        from truthound.stores.versioning import (
            VersionedStore,
            VersioningConfig,
            IncrementalStrategy,
        )
        from truthound.stores.caching import CachedStore, LRUCache, CacheMode
        from truthound.stores.observability import ObservableStore, ObservabilityConfig, AuditConfig, MetricsConfig

        # Build layered store
        base = get_store("memory")

        # Add caching
        cache = LRUCache(max_size=100)
        cached = CachedStore(base, cache, mode=CacheMode.WRITE_THROUGH)

        # Add versioning
        version_config = VersioningConfig(strategy=IncrementalStrategy())
        versioned = VersionedStore(cached, version_config)

        # Add observability
        obs_config = ObservabilityConfig(
            audit=AuditConfig(enabled=True),
            metrics=MetricsConfig(enabled=True),
        )
        store = ObservableStore(versioned, obs_config)

        # Use composed store
        result = ValidationResult(
            run_id="composed_test_001",
            run_time=datetime.now(timezone.utc),
            data_asset="composed_data.csv",
            status=ResultStatus.SUCCESS,
            results=[],
            statistics=ResultStatistics(total_validators=5, passed_validators=5),
            tags={},
            metadata={},
        )

        run_id = store.save(result)
        retrieved = store.get(run_id)

        assert retrieved is not None, "Should retrieve through composed store"

        # Check versioning worked
        history = versioned.get_version_history("composed_data.csv")
        assert len(history) >= 1, "Should have version history"

        # Check observability worked
        events = store.get_audit_events()
        assert len(events) >= 2, "Should have audit events"

        store.close()
        return True


def run_all_tests():
    """Run all storage feature tests."""
    print_header("TRUTHOUND STORAGE FEATURE TESTS")
    print(f"Started at: {datetime.now()}")

    tests = StorageFeatureTests()

    try:
        # 1. Factory and Backend Tests
        print_header("1. Factory and Backend Tests")
        tests.run_test("List available backends", tests.test_factory_list_backends)
        tests.run_test("Filesystem store basic operations", tests.test_filesystem_store_basic)
        tests.run_test("Memory store basic operations", tests.test_memory_store_basic)
        tests.run_test("Store query functionality", tests.test_store_query)

        # 2. Versioning Tests
        print_header("2. Versioning System Tests")
        tests.run_test("Incremental versioning strategy", tests.test_versioning_incremental)
        tests.run_test("Semantic versioning strategy", tests.test_versioning_semantic)
        tests.run_test("Timestamp versioning strategy", tests.test_versioning_timestamp)
        tests.run_test("Git-like versioning strategy", tests.test_versioning_gitlike)

        # 3. Retention Tests
        print_header("3. Retention Policy Tests")
        tests.run_test("Time-based retention policy", tests.test_retention_time_based)
        tests.run_test("Count-based retention policy", tests.test_retention_count_based)
        tests.run_test("Size-based retention policy", tests.test_retention_size_based)
        tests.run_test("Composite retention policy", tests.test_retention_composite)

        # 4. Tiering Tests
        print_header("4. Tiered Storage Tests")
        tests.run_test("Basic tiering functionality", tests.test_tiering_basic)
        tests.run_test("Age-based tier policy", tests.test_tiering_age_based_policy)
        tests.run_test("Access-based tier policy", tests.test_tiering_access_based_policy)

        # 5. Caching Tests
        print_header("5. Caching Layer Tests")
        tests.run_test("In-memory cache", tests.test_cache_inmemory)
        tests.run_test("LRU cache eviction", tests.test_cache_lru)
        tests.run_test("LFU cache eviction", tests.test_cache_lfu)
        tests.run_test("TTL cache expiration", tests.test_cache_ttl)
        tests.run_test("CachedStore wrapper", tests.test_cached_store)

        # 6. Batch Write Tests
        print_header("6. Batch Write Optimization Tests")
        tests.run_test("Batch buffer", tests.test_batch_buffer)
        tests.run_test("Memory-aware buffer", tests.test_memory_aware_buffer)
        tests.run_test("Batch writer", tests.test_batch_writer)
        tests.run_test("Async batch writer", tests.test_async_batch_writer)

        # 7. Backpressure Tests
        print_header("7. Backpressure Management Tests")
        tests.run_test("Memory-based backpressure", tests.test_backpressure_memory)
        tests.run_test("Queue-depth backpressure", tests.test_backpressure_queue_depth)
        tests.run_test("Token bucket rate limiter", tests.test_backpressure_token_bucket)
        tests.run_test("Circuit breaker pattern", tests.test_circuit_breaker)

        # 8. Replication Tests
        print_header("8. Cross-Region Replication Tests")
        tests.run_test("Basic replication setup", tests.test_replication_basic)
        tests.run_test("Synchronous replication", tests.test_replication_sync)

        # 9. Migration Tests
        print_header("9. Schema Migration Tests")
        tests.run_test("Migration registry", tests.test_migration_registry)
        tests.run_test("Migration execution", tests.test_migration_execution)

        # 10. Encryption Tests
        print_header("10. Encryption Module Tests")
        tests.run_test("Basic encryption", tests.test_encryption_basic)
        tests.run_test("Multiple encryption algorithms", tests.test_encryption_algorithms)
        tests.run_test("Encryption pipeline", tests.test_encryption_pipeline)

        # 11. Compression Tests
        print_header("11. Compression Module Tests")
        tests.run_test("Basic compression", tests.test_compression_basic)
        tests.run_test("Multiple compression algorithms", tests.test_compression_algorithms)
        tests.run_test("Compression pipeline", tests.test_compression_pipeline)

        # 12. Observability Tests
        print_header("12. Observability Layer Tests")
        tests.run_test("Audit logging", tests.test_audit_logging)
        tests.run_test("Metrics collection", tests.test_metrics_collection)
        tests.run_test("Observable store wrapper", tests.test_observable_store)
        tests.run_test("Distributed tracing", tests.test_tracing)

        # 13. Streaming Tests
        print_header("13. Streaming Storage Tests")
        tests.run_test("Streaming filesystem store", tests.test_streaming_filesystem)

        # 14. Concurrency Tests
        print_header("14. Concurrency Tests")
        tests.run_test("Concurrent filesystem store", tests.test_concurrent_filesystem)

        # 15. Store Composition Tests
        print_header("15. Store Composition Tests")
        tests.run_test("Store feature composition", tests.test_store_composition)

    finally:
        tests.cleanup()

    # Print summary
    print_header("TEST SUMMARY")
    total = tests.passed + tests.failed
    print(f"  Total Tests: {total}")
    print(f"  Passed:      {tests.passed} ({100*tests.passed/total:.1f}%)")
    print(f"  Failed:      {tests.failed} ({100*tests.failed/total:.1f}%)")

    if tests.errors:
        print(f"\n  Errors:")
        for error in tests.errors:
            print(f"    - {error}")

    print(f"\nFinished at: {datetime.now()}")

    return tests.failed == 0


if __name__ == "__main__":
    success = run_all_tests()
    sys.exit(0 if success else 1)
