"""
Comprehensive tests for streaming functionality.

Tests cover:
1. Large file streaming (>100MB)
2. Memory usage profiling
3. Streaming error recovery
4. Partial chunk handling
5. Buffer size optimization
6. Progress callbacks
"""

import concurrent.futures
import mmap
import shutil
import tempfile
import threading
import time
from collections.abc import Callable, Iterator
from pathlib import Path

import psutil
import pytest

from chunker import chunk_file
from chunker.exceptions import LanguageNotFoundError
from chunker.streaming import (
    StreamingChunker,
    chunk_file_streaming,
    compute_file_hash,
    get_file_metadata,
)
from chunker.types import CodeChunk


def generate_large_python_code(num_functions: int = 1000) -> str:
    """Generate a large Python file with many functions."""
    code_parts = ['"""Large auto-generated Python file for testing.\n']
    code_parts.append(
        """This file is automatically generated for testing streaming functionality.
"""
        * 500,
    )
    code_parts.append('"""\n\n')
    imports = [
        "import math",
        "import sys",
        "import os",
        "import json",
        "import datetime",
        "import collections",
        "import itertools",
        "import functools",
        "import re",
        "import typing",
        "import pathlib",
        "import tempfile",
        "import shutil",
    ]
    code_parts.extend(imp + "\n" for imp in imports)
    code_parts.append("\n# " + "=" * 100 + "\n\n")
    for i in range(num_functions // 10):
        code_parts.append(f"class TestClass{i}:\n")
        code_parts.append(f'    """Test class number {i}.\n')
        code_parts.append("    \n")
        code_parts.append(
            """    This is a test class with extensive documentation to increase file size.
"""
            * 10,
        )
        code_parts.append("    It contains multiple methods for testing purposes.\n")
        code_parts.append('    """\n\n')
        code_parts.extend(
            f"    class_var_{v} = '" + "x" * 200 + "'  # Long string variable\n"
            for v in range(5)
        )
        code_parts.append("\n")
        for j in range(10):
            func_num = i * 10 + j
            code_parts.append(
                f"    def method_{func_num}(self, x, y, z=None, *args, **kwargs):\n",
            )
            code_parts.append(f'        """Method {func_num} in class {i}.\n')
            code_parts.append("        \n")
            code_parts.append(
                """        Detailed documentation for this method with multiple lines.
"""
                * 5,
            )
            code_parts.append('        """\n')
            code_parts.append(f"        # This is method number {func_num}\n")
            code_parts.append("        # " + "=" * 50 + "\n")
            code_parts.append("        data = {\n")
            code_parts.extend(
                f"            'key_{k}': '" + "value" * 20 + "',\n" for k in range(10)
            )
            code_parts.append("        }\n")
            code_parts.append(f"        result = x * {func_num} + y * {func_num % 7}\n")
            code_parts.append("        long_string = '" + "test" * 100 + "'\n")
            for k in range(5):
                code_parts.append(f"        if result > {k * 10}:\n")
                code_parts.append(f"            result = result * {k + 1}\n")
                code_parts.append(f"            # Processing step {k}\n")
            code_parts.append(f"        return result + {func_num}\n\n")
    for i in range(num_functions % 10):
        code_parts.append(
            f"def standalone_function_{i}(param1, param2, param3=None):\n",
        )
        code_parts.append(f'    """Standalone function {i}.\n')
        code_parts.append("    \n")
        code_parts.append("    This function does important processing.\n" * 5)
        code_parts.append('    """\n')
        code_parts.append("    # Function implementation\n")
        code_parts.append(
            "    large_list = [" + ", ".join(str(x) for x in range(100)) + "]\n",
        )
        code_parts.append(f"    return param1 + param2 * {i}\n\n")
    return "".join(code_parts)


class MemoryMonitor:
    """Monitor memory usage during streaming operations."""

    def __init__(self):
        self.process = psutil.Process()
        self.initial_memory = 0
        self.peak_memory = 0
        self.samples = []
        self.monitoring = False
        self._lock = threading.Lock()

    def start(self):
        """Start monitoring memory usage."""
        self.initial_memory = self.process.memory_info().rss
        self.peak_memory = self.initial_memory
        self.samples = []
        self.monitoring = True
        self._monitor_thread = threading.Thread(target=self._monitor_loop)
        self._monitor_thread.daemon = True
        self._monitor_thread.start()

    def stop(self):
        """Stop monitoring and return statistics."""
        self.monitoring = False
        self._monitor_thread.join(timeout=1)
        with self._lock:
            return {
                "initial_mb": self.initial_memory / (1024 * 1024),
                "peak_mb": self.peak_memory / (1024 * 1024),
                "increase_mb": (self.peak_memory - self.initial_memory) / (1024 * 1024),
                "num_samples": len(self.samples),
                "average_mb": (
                    sum(self.samples) / len(self.samples) / (1024 * 1024)
                    if self.samples
                    else 0
                ),
            }

    def _monitor_loop(self):
        """Monitor memory usage in a loop."""
        while self.monitoring:
            current_memory = self.process.memory_info().rss
            with self._lock:
                self.samples.append(current_memory)
                self.peak_memory = max(self.peak_memory, current_memory)
            time.sleep(0.1)


@pytest.fixture
def large_python_file():
    """Create a large temporary Python file (>100MB)."""
    with tempfile.NamedTemporaryFile(
        encoding="utf-8",
        mode="w",
        suffix=".py",
        delete=False,
    ) as f:
        num_functions = 10000
        large_code = generate_large_python_code(num_functions=num_functions)
        while len(large_code.encode("utf-8")) < 100 * 1024 * 1024:
            num_functions *= 2
            large_code = generate_large_python_code(
                num_functions=num_functions,
            )
        f.write(large_code)
        temp_path = Path(f.name)
    actual_size_mb = temp_path.stat().st_size / (1024 * 1024)
    assert (
        actual_size_mb > 100
    ), f"Generated file is only {actual_size_mb:.1f}MB, need 100MB+"
    yield temp_path
    temp_path.unlink()


@pytest.fixture
def medium_python_file():
    """Create a medium-sized temporary Python file (~10MB)."""
    with tempfile.NamedTemporaryFile(
        encoding="utf-8",
        mode="w",
        suffix=".py",
        delete=False,
    ) as f:
        medium_code = generate_large_python_code(num_functions=5000)
        f.write(medium_code)
        temp_path = Path(f.name)
    yield temp_path
    temp_path.unlink()


@pytest.fixture
def corrupted_python_file():
    """Create a file with invalid UTF-8 sequences."""
    with tempfile.NamedTemporaryFile(
        mode="wb",
        suffix=".py",
        delete=False,
    ) as f:
        f.write(b"def valid_function():\n    pass\n\n")
        f.write(
            b"def corrupted_function():\n    # Invalid UTF-8: \xff\xfe\n    pass\n\n",
        )
        f.write(b"def another_valid_function():\n    return 42\n")
        temp_path = Path(f.name)
    yield temp_path
    temp_path.unlink()


class TestStreamingLargeFiles:
    """Test streaming functionality with large files."""

    @classmethod
    def test_large_file_streaming(cls, large_python_file):
        """Test streaming a large file (>100MB) without loading it entirely into memory."""
        monitor = MemoryMonitor()
        monitor.start()
        chunk_count = 0
        chunker = StreamingChunker("python")
        for chunk in chunker.chunk_file_streaming(large_python_file):
            chunk_count += 1
            assert isinstance(chunk, CodeChunk)
            assert chunk.language == "python"
            assert chunk.content
            assert chunk.node_type in {
                "function_definition",
                "class_definition",
                "method_definition",
            }
        memory_stats = monitor.stop()
        assert chunk_count > 0
        file_size_mb = large_python_file.stat().st_size / (1024 * 1024)
        max_allowed_mb = file_size_mb * 30
        assert (
            memory_stats["increase_mb"] < max_allowed_mb
        ), f"Memory increase too high: {memory_stats['increase_mb']}MB for {file_size_mb}MB file"
        file_size_mb = large_python_file.stat().st_size / (1024 * 1024)
        assert file_size_mb > 100, f"Test file too small: {file_size_mb}MB"

    @classmethod
    def test_streaming_vs_regular_memory_usage(cls, medium_python_file):
        """Compare memory usage between streaming and regular chunking."""
        monitor_regular = MemoryMonitor()
        monitor_regular.start()
        regular_chunks = chunk_file(medium_python_file, "python")
        regular_stats = monitor_regular.stop()
        monitor_streaming = MemoryMonitor()
        monitor_streaming.start()
        streaming_chunks = list(chunk_file_streaming(medium_python_file, "python"))
        streaming_stats = monitor_streaming.stop()
        assert len(regular_chunks) == len(streaming_chunks)
        assert streaming_stats["peak_mb"] <= regular_stats["peak_mb"] * 2.0

    @classmethod
    def test_streaming_retrieval_metadata_matches_regular(cls, medium_python_file):
        """Streaming chunks should include the same retrieval metadata contract."""
        regular_chunks = chunk_file(
            medium_python_file,
            "python",
            include_retrieval_metadata=True,
        )
        streaming_chunks = list(
            chunk_file_streaming(
                medium_python_file,
                "python",
                include_retrieval_metadata=True,
            )
        )

        regular_by_route = {
            tuple(chunk.qualified_route): chunk for chunk in regular_chunks
        }
        streaming_by_route = {
            tuple(chunk.qualified_route): chunk for chunk in streaming_chunks
        }

        shared_routes = set(regular_by_route) & set(streaming_by_route)
        assert shared_routes

        sample_route = next(iter(shared_routes))
        regular_chunk = regular_by_route[sample_route]
        streaming_chunk = streaming_by_route[sample_route]

        assert streaming_chunk.metadata["kind"] == regular_chunk.metadata["kind"]
        assert streaming_chunk.metadata["symbol"] == regular_chunk.metadata["symbol"]
        assert (
            streaming_chunk.metadata["qualified_name"]
            == regular_chunk.metadata["qualified_name"]
        )


class TestMemoryEfficiency:
    """Test memory efficiency and profiling."""

    @classmethod
    def test_memory_mapped_file_access(cls, medium_python_file):
        """Test that memory-mapped file access is working correctly."""
        StreamingChunker("python")
        with (
            Path(medium_python_file).open("rb") as f,
            mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mmap_data,
        ):
            assert len(mmap_data) > 0
            first_100_bytes = mmap_data[:100]
            assert len(first_100_bytes) == 100

        with (
            Path(medium_python_file).open(
                "rb",
            ) as f,
            mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mmap_data,
        ):
            # Test direct access
            assert len(mmap_data) > 0

            # Test slicing
            first_100_bytes = mmap_data[:100]
            assert len(first_100_bytes) == 100

    def test_progressive_memory_usage(self, medium_python_file):
        """Test that memory usage doesn't grow linearly with chunks processed."""
        monitor = MemoryMonitor()
        monitor.start()
        memory_checkpoints = []
        chunk_count = 0
        for _chunk in chunk_file_streaming(medium_python_file, "python"):
            chunk_count += 1
            if chunk_count % 100 == 0:
                current_memory = psutil.Process().memory_info().rss / (1024 * 1024)
                memory_checkpoints.append(current_memory)
        monitor.stop()
        if len(memory_checkpoints) > 2:
            initial_checkpoint = memory_checkpoints[0]
            final_checkpoint = memory_checkpoints[-1]
            memory_growth = final_checkpoint - initial_checkpoint
            assert (
                memory_growth < 200
            ), f"Memory grew by {memory_growth}MB during streaming"


class TestStreamingErrorRecovery:
    """Test error handling and recovery in streaming operations."""

    @staticmethod
    def test_corrupted_file_handling(corrupted_python_file):
        """Test handling of files with invalid UTF-8 sequences."""
        chunks = list(chunk_file_streaming(corrupted_python_file, "python"))
        assert len(chunks) > 0
        function_names = [
            chunk.content.split("(")[0].split()[-1]
            for chunk in chunks
            if chunk.node_type == "function_definition"
        ]
        assert "valid_function" in function_names
        assert "another_valid_function" in function_names

    @staticmethod
    def test_file_not_found_error():
        """Test handling of non-existent files."""
        with pytest.raises(FileNotFoundError):
            list(chunk_file_streaming("/nonexistent/file.py", "python"))

    @staticmethod
    def test_unsupported_language_error(medium_python_file):
        """Test handling of unsupported languages."""
        with pytest.raises(LanguageNotFoundError):
            list(chunk_file_streaming(medium_python_file, "unsupported_lang"))

    @classmethod
    def test_permission_error_handling(cls):
        """Test handling of permission errors."""
        with tempfile.NamedTemporaryFile(
            encoding="utf-8",
            mode="w",
            suffix=".py",
            delete=False,
        ) as f:
            f.write("def test(): pass")
            temp_path = Path(f.name)
        try:
            Path(temp_path).chmod(0)
            with pytest.raises(PermissionError):
                list(chunk_file_streaming(temp_path, "python"))
        finally:
            Path(temp_path).chmod(0o644)
            temp_path.unlink()


class TestPartialChunkHandling:
    """Test handling of partial chunks and boundaries."""

    @classmethod
    def test_chunk_boundary_integrity(cls, medium_python_file):
        """Test that chunk boundaries are properly maintained."""
        chunks = list(chunk_file_streaming(medium_python_file, "python"))
        for chunk in chunks:
            with Path(medium_python_file).open("rb") as f:
                f.seek(chunk.byte_start)
                expected_content = f.read(chunk.byte_end - chunk.byte_start)
                assert (
                    chunk.content.encode(
                        "utf-8",
                        errors="replace",
                    )
                    == expected_content
                )

    @classmethod
    def test_nested_chunk_handling(cls):
        """Test handling of nested code structures."""
        nested_code = """
class OuterClass:
    class InnerClass:
        def inner_method(self):
            def nested_function():
                return 42
            return nested_function()

    def outer_method(self):
        return self.InnerClass()
"""
        with tempfile.NamedTemporaryFile(
            encoding="utf-8",
            mode="w",
            suffix=".py",
            delete=False,
        ) as f:
            f.write(nested_code)
            temp_path = Path(f.name)
        try:
            chunks = list(chunk_file_streaming(temp_path, "python"))
            chunk_types = [c.node_type for c in chunks]
            assert "class_definition" in chunk_types
            assert "function_definition" in chunk_types
            method_chunks = [
                c
                for c in chunks
                if "inner_method" in c.content or "outer_method" in c.content
            ]
            assert len(method_chunks) >= 2
        finally:
            temp_path.unlink()


class TestBufferOptimization:
    """Test buffer size optimization and performance."""

    @staticmethod
    def test_file_hash_computation_performance(large_python_file):
        """Test efficient file hash computation."""
        start_time = time.time()
        hash1 = compute_file_hash(large_python_file)
        default_time = time.time() - start_time
        start_time = time.time()
        hash2 = compute_file_hash(large_python_file, chunk_size=1024 * 1024)
        large_chunk_time = time.time() - start_time
        assert hash1 == hash2
        assert large_chunk_time <= default_time * 1.1

    @staticmethod
    def test_streaming_performance_consistency(medium_python_file):
        """Test that streaming performance is consistent across runs."""
        # Warmup run (not timed) to initialize caches and JIT
        list(chunk_file_streaming(medium_python_file, "python"))

        times = []
        for _ in range(3):
            start_time = time.time()
            list(chunk_file_streaming(medium_python_file, "python"))
            elapsed = time.time() - start_time
            times.append(elapsed)
        avg_time = sum(times) / len(times)
        variance = sum((t - avg_time) ** 2 for t in times) / len(times)
        # Increased tolerance from 0.01 to 0.05 (5% variance is acceptable)
        assert variance < 0.05, f"High variance in streaming times: {variance}"


class TestProgressCallbacks:
    """Test progress callback functionality."""

    @staticmethod
    def test_progress_callback_integration(medium_python_file):
        """Test integration of progress callbacks with streaming."""
        progress_calls = []

        def progress_callback(current: int, total: int, chunk: CodeChunk | None = None):
            progress_calls.append(
                {"current": current, "total": total, "has_chunk": chunk is not None},
            )

        class ProgressStreamingChunker(StreamingChunker):
            def __init__(
                self,
                language: str,
                progress_callback: Callable | None = None,
            ):
                super().__init__(language)
                self.progress_callback = progress_callback
                self._chunk_count = 0

            def chunk_file_streaming(
                self,
                path: Path,
                include_retrieval_metadata: bool = False,
            ) -> Iterator[CodeChunk]:
                file_size = path.stat().st_size
                for chunk in super().chunk_file_streaming(
                    path,
                    include_retrieval_metadata=include_retrieval_metadata,
                ):
                    self._chunk_count += 1
                    if self.progress_callback:
                        self.progress_callback(
                            chunk.byte_end,
                            file_size,
                            chunk,
                        )
                    yield chunk

        chunker = ProgressStreamingChunker("python", progress_callback)
        chunks = list(chunker.chunk_file_streaming(medium_python_file))
        assert len(progress_calls) > 0
        assert len(progress_calls) == len(chunks)

    @staticmethod
    def test_cancellable_streaming(large_python_file):
        """Test ability to cancel streaming operation."""

        class CancellableStreamingChunker(StreamingChunker):
            def __init__(self, language: str):
                super().__init__(language)
                self.cancelled = False

            def chunk_file_streaming(
                self,
                path: Path,
                include_retrieval_metadata: bool = False,
            ) -> Iterator[CodeChunk]:
                for chunk in super().chunk_file_streaming(
                    path,
                    include_retrieval_metadata=include_retrieval_metadata,
                ):
                    if self.cancelled:
                        break
                    yield chunk

        chunker = CancellableStreamingChunker("python")
        chunks_processed = 0
        for _chunk in chunker.chunk_file_streaming(large_python_file):
            chunks_processed += 1
            if chunks_processed >= 10:
                chunker.cancelled = True
        assert chunks_processed <= 11


class TestFileMetadata:
    """Test file metadata functionality."""

    @staticmethod
    def test_get_file_metadata(medium_python_file):
        """Test file metadata extraction."""
        metadata = get_file_metadata(medium_python_file)
        assert metadata.path == str(medium_python_file)
        assert metadata.size > 0
        assert len(metadata.hash) == 64
        assert metadata.mtime > 0

    @classmethod
    def test_metadata_caching_validity(cls, medium_python_file):
        """Test that metadata can be used for cache validation."""
        metadata1 = get_file_metadata(medium_python_file)
        metadata2 = get_file_metadata(medium_python_file)
        assert metadata1.hash == metadata2.hash
        assert metadata1.size == metadata2.size
        time.sleep(0.01)
        with Path(medium_python_file).open("a", encoding="utf-8") as f:
            f.write("\n# Modified\n")
        metadata3 = get_file_metadata(medium_python_file)
        assert metadata3.hash != metadata1.hash
        assert metadata3.size > metadata1.size
        assert metadata3.mtime >= metadata1.mtime


class TestStreamingEdgeCases:
    """Test edge cases and special scenarios."""

    @classmethod
    def test_empty_file_handling(cls):
        """Test handling of empty files."""
        with tempfile.NamedTemporaryFile(
            encoding="utf-8",
            mode="w",
            suffix=".py",
            delete=False,
        ) as f:
            temp_path = Path(f.name)
        try:
            chunks = list(chunk_file_streaming(temp_path, "python"))
            assert len(chunks) == 0
        finally:
            temp_path.unlink()

    @classmethod
    def test_single_line_file(cls):
        """Test handling of single-line files."""
        with tempfile.NamedTemporaryFile(
            encoding="utf-8",
            mode="w",
            suffix=".py",
            delete=False,
        ) as f:
            f.write("def oneliner(): return 42")
            temp_path = Path(f.name)
        try:
            chunks = list(chunk_file_streaming(temp_path, "python"))
            assert len(chunks) == 1
            assert chunks[0].node_type == "function_definition"
            assert chunks[0].start_line == 1
            assert chunks[0].end_line == 1
        finally:
            temp_path.unlink()

    @classmethod
    def test_file_with_no_chunks(cls):
        """Test handling of files with no chunkable content."""
        with tempfile.NamedTemporaryFile(
            encoding="utf-8",
            mode="w",
            suffix=".py",
            delete=False,
        ) as f:
            f.write(
                "# Just comments\n# No functions or classes\nimport os\nVARIABLE = 42\n",
            )
            temp_path = Path(f.name)
        try:
            chunks = list(chunk_file_streaming(temp_path, "python"))
            assert len(chunks) == 0
        finally:
            temp_path.unlink()

    @classmethod
    @pytest.mark.parametrize("encoding", ["utf-8", "latin-1", "utf-16"])
    def test_different_encodings(cls, encoding):
        """Test handling of files with different encodings."""
        content = 'def test_encoding():\n    return "Hello, World!"\n'
        with tempfile.NamedTemporaryFile(
            mode="w",
            suffix=".py",
            delete=False,
            encoding=encoding,
        ) as f:
            try:
                f.write(content)
            except UnicodeEncodeError:
                pytest.skip(f"Cannot encode test content with {encoding}")
            temp_path = Path(f.name)
        try:
            chunks = list(chunk_file_streaming(temp_path, "python"))
            if encoding == "utf-16":
                assert isinstance(chunks, list)
            else:
                assert len(chunks) >= 1
        finally:
            temp_path.unlink()


class TestConcurrentStreaming:
    """Test concurrent streaming operations."""

    @classmethod
    def test_multiple_files_concurrent_streaming(
        cls,
        temp_directory_with_files,
    ):
        """Test streaming multiple files concurrently."""
        temp_dir = Path(tempfile.mkdtemp())
        files = []
        try:
            for i in range(5):
                file_path = temp_dir / f"concurrent_test_{i}.py"
                file_path.write_text(generate_large_python_code(num_functions=100))
                files.append(file_path)
            all_chunks = {}

            def process_file(file_path):
                return file_path, list(chunk_file_streaming(file_path, "python"))

            with concurrent.futures.ThreadPoolExecutor(
                max_workers=3,
            ) as executor:
                future_to_file = {executor.submit(process_file, f): f for f in files}
                for future in concurrent.futures.as_completed(future_to_file):
                    file_path, chunks = future.result()
                    all_chunks[file_path] = chunks
            assert len(all_chunks) == len(files)
            for file_path, chunks in all_chunks.items():
                assert len(chunks) > 0
        finally:
            shutil.rmtree(temp_dir)

    @classmethod
    def test_thread_safety(cls, medium_python_file):
        """Test that StreamingChunker is thread-safe."""
        chunker = StreamingChunker("python")
        results = []

        def stream_chunks():
            return list(chunker.chunk_file_streaming(medium_python_file))

        with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
            futures = [executor.submit(stream_chunks) for _ in range(3)]
            results.extend(
                future.result() for future in concurrent.futures.as_completed(futures)
            )
        assert len(results) == 3
        first_result = results[0]
        for result in results[1:]:
            assert len(result) == len(first_result)
            assert [c.chunk_id for c in result] == [c.chunk_id for c in first_result]


@pytest.fixture
def temp_directory_with_files():
    """Create a temporary directory with multiple Python files."""
    temp_dir = Path(tempfile.mkdtemp())
    for i in range(5):
        file_path = temp_dir / f"test_file_{i}.py"
        file_path.write_text(generate_large_python_code(num_functions=50))
    yield temp_dir
    shutil.rmtree(temp_dir)
