#!/usr/bin/env python
"""Comprehensive tests for all node types in the soak pipeline system."""

from pathlib import Path

import pytest
from struckdown import LLMCredentials

from soak.specs import load_template_bundle

# Test data path
TEST_DATA = Path("data-private/chainstorecoffee.txt")


@pytest.mark.anyio
async def test_split_node():
    """Test Split node with different split_unit options."""
    pipeline = load_template_bundle(Path("tests/pipelines/test_split.soak"))
    pipeline.config.document_paths = [str(TEST_DATA)]
    pipeline.config.llm_credentials = LLMCredentials()

    result, error = await pipeline.run()
    assert error is None, f"Pipeline failed: {error}"

    # Test each split type produced output
    for node in pipeline.nodes:
        assert node.output is not None, f"Node {node.name} has no output"
        assert len(node.output) > 0, f"Node {node.name} produced empty output"

        # Verify all chunks have content
        for chunk in node.output:
            content = chunk.content if hasattr(chunk, "content") else chunk
            assert len(content) > 0, f"Node {node.name} has empty chunk"

        # Verify provenance tracking
        if hasattr(node.output[0], "source_id"):
            assert (
                node.name in node.output[0].source_id
            ), f"Node name not in source_id for {node.name}"

    # Test specific behaviors
    tokens_node = [n for n in pipeline.nodes if n.name == "split_tokens"][0]
    words_node = [n for n in pipeline.nodes if n.name == "split_words"][0]
    sentences_node = [n for n in pipeline.nodes if n.name == "split_sentences"][0]
    paragraphs_node = [n for n in pipeline.nodes if n.name == "split_paragraphs"][0]

    # Document is long enough to require multiple chunks
    assert len(tokens_node.output) > 1, "tokens split should produce multiple chunks"
    assert len(words_node.output) > 1, "words split should produce multiple chunks"
    assert (
        len(sentences_node.output) > 1
    ), "sentences split should produce multiple chunks"
    assert (
        len(paragraphs_node.output) > 1
    ), "paragraphs split should produce multiple chunks"

    # Test overlap functionality
    assert tokens_node.overlap == 50, "tokens split should have overlap of 50"
    assert words_node.overlap == 30, "words split should have overlap of 30"


@pytest.mark.anyio
async def test_map_node():
    """Test Map node processes all items in parallel."""
    pipeline = load_template_bundle(Path("tests/pipelines/test_map.soak"))
    pipeline.config.document_paths = [str(TEST_DATA)]
    pipeline.config.llm_credentials = LLMCredentials()

    result, error = await pipeline.run()
    assert error is None, f"Pipeline failed: {error}"

    chunks_node = [n for n in pipeline.nodes if n.name == "chunks"][0]
    map_node = [n for n in pipeline.nodes if n.name == "map_summarize"][0]

    # Map should process each chunk
    assert len(map_node.output) == len(
        chunks_node.output
    ), "Map output count should equal input count"

    # All outputs should have content
    for item in map_node.output:
        content = item.content if hasattr(item, "content") else str(item)
        assert len(content) > 0, "Map produced empty output"


@pytest.mark.anyio
async def test_classifier_node_with_agreement():
    """Test Classifier node with multiple models and agreement calculation."""
    pipeline = load_template_bundle(Path("tests/pipelines/test_classifier.soak"))
    pipeline.config.document_paths = [str(TEST_DATA)]
    pipeline.config.llm_credentials = LLMCredentials()

    result, error = await pipeline.run()
    assert error is None, f"Pipeline failed: {error}"

    classifier_node = [n for n in pipeline.nodes if n.name == "classify"][0]

    # Should have output
    assert classifier_node.output is not None
    assert len(classifier_node.output) > 0

    # Check that multiple models were used
    assert classifier_node.model_names is not None
    assert len(classifier_node.model_names) == 2

    # Check agreement fields are set
    assert classifier_node.agreement_fields is not None
    assert "topic" in classifier_node.agreement_fields
    assert "sentiment" in classifier_node.agreement_fields

    # Check provenance - each item should have model info
    for item in classifier_node.output:
        if hasattr(item, "metadata"):
            assert "model_name" in item.metadata or "source_id" in item.metadata

    # Export and check agreement stats (if calculated)
    import tempfile

    with tempfile.TemporaryDirectory() as tmpdir:
        folder = Path(tmpdir)
        classifier_node.export(folder)

        # Check CSV files were created for each model
        csv_files = list(folder.glob("*.csv"))
        assert len(csv_files) >= 2, "Should have CSV files for multiple models"

        # Check for agreement stats files
        agreement_stats_csv = folder / "agreement_stats.csv"
        agreement_stats_json = folder / "agreement_stats.json"

        # Agreement stats should exist if we have multiple models
        if classifier_node._agreement_stats:
            assert agreement_stats_csv.exists(), "Agreement stats CSV should exist"
            assert agreement_stats_json.exists(), "Agreement stats JSON should exist"


@pytest.mark.anyio
async def test_filter_node():
    """Test Filter node includes and excludes items correctly."""
    pipeline = load_template_bundle(Path("tests/pipelines/test_filter.soak"))
    pipeline.config.document_paths = [str(TEST_DATA)]
    pipeline.config.llm_credentials = LLMCredentials()

    result, error = await pipeline.run()
    assert error is None, f"Pipeline failed: {error}"

    chunks_node = [n for n in pipeline.nodes if n.name == "chunks"][0]
    filter_node = [n for n in pipeline.nodes if n.name == "filter_about_coffee"][0]

    # Some items should pass filter (coffee article)
    assert filter_node.output is not None
    assert len(filter_node.output) > 0, "Filter should pass some items about coffee"

    # Some items should be excluded
    assert hasattr(filter_node, "_excluded_items")

    # Total should equal input
    total_processed = len(filter_node.output) + len(filter_node._excluded_items or [])
    assert total_processed == len(
        chunks_node.output
    ), "Filter should process all input items"


@pytest.mark.anyio
async def test_filter_preserves_metadata():
    """Test that Filter node preserves TrackedItem metadata for excluded items."""
    from soak.models.base import TrackedItem
    from soak.models.dag import DAG, DAGConfig
    from soak.models.nodes.filter import Filter

    # create test items with metadata
    items = [
        TrackedItem(
            content="This is about coffee",
            id="doc1__chunk__0",
            sources=["doc1"],
            metadata={"original_path": "test.txt", "index": 0, "category": "food"},
        ),
        TrackedItem(
            content="This is about tea",
            id="doc1__chunk__1",
            sources=["doc1"],
            metadata={"original_path": "test.txt", "index": 1, "category": "food"},
        ),
    ]

    # create minimal DAG
    config = DAGConfig(llm_credentials=LLMCredentials())
    dag = DAG(name="test_filter", config=config)

    # create filter node (simple mode)
    filter_node = Filter(
        name="test_filter",
        inputs=[],
        expression="'coffee' in input",
        mode="simple",
        dag=dag,
    )

    # process items
    result = await filter_node.process_items(items)

    # should still have 2 items (included + excluded with "...")
    assert len(result) == 2, "Filter should keep both included and excluded items"

    # first item should be included (has coffee)
    included_item = result[0]
    assert included_item.content == "This is about coffee"
    assert included_item.metadata.get("original_path") == "test.txt"
    assert included_item.metadata.get("index") == 0
    assert included_item.metadata.get("category") == "food"
    assert "filtered" not in included_item.metadata

    # second item should be excluded (no coffee)
    excluded_item = result[1]
    assert excluded_item.content == " ... ", "Excluded item should have omitted_text"
    assert (
        excluded_item.metadata.get("filtered") == True
    ), "Excluded item should have filtered=True"
    assert (
        excluded_item.metadata.get("original_path") == "test.txt"
    ), "Excluded item should preserve original_path"
    assert (
        excluded_item.metadata.get("index") == 1
    ), "Excluded item should preserve index"
    assert (
        excluded_item.metadata.get("category") == "food"
    ), "Excluded item should preserve category"

    # verify IDs are preserved
    assert excluded_item.id == "doc1__chunk__1", "Excluded item should preserve ID"
    assert excluded_item.sources == ["doc1"], "Excluded item should preserve sources"


@pytest.mark.anyio
async def test_transform_node_multi_slot():
    """Test Transform node with multi-slot struckdown prompts."""
    pipeline = load_template_bundle(Path("tests/pipelines/test_transform.soak"))
    pipeline.config.document_paths = [str(TEST_DATA)]
    pipeline.config.llm_credentials = LLMCredentials()

    result, error = await pipeline.run()
    assert error is None, f"Pipeline failed: {error}"

    transform_node = [n for n in pipeline.nodes if n.name == "multi_slot_transform"][0]

    # Should have output
    assert transform_node.output is not None

    # Transform output is a ChatterResult with structured data
    output = transform_node.output
    assert hasattr(
        output, "response"
    ), "Transform output should have response attribute"

    # Response can be a string or dict depending on the struckdown output
    response_str = str(output.response)
    assert len(response_str) > 0, "Transform should produce non-empty output"


@pytest.mark.anyio
async def test_reduce_node():
    """Test Reduce node aggregates multiple inputs."""
    pipeline = load_template_bundle(Path("tests/pipelines/test_reduce.soak"))
    pipeline.config.document_paths = [str(TEST_DATA)]
    pipeline.config.llm_credentials = LLMCredentials()

    result, error = await pipeline.run()
    assert error is None, f"Pipeline failed: {error}"

    summaries_node = [n for n in pipeline.nodes if n.name == "summaries"][0]
    reduce_node = [n for n in pipeline.nodes if n.name == "combined"][0]

    # Reduce should aggregate multiple inputs into single output (string)
    assert len(summaries_node.output) > 1, "Should have multiple summaries"
    assert reduce_node.output is not None

    # Reduce concatenates into a single string
    assert isinstance(reduce_node.output, str), "Reduce output should be a string"
    assert len(reduce_node.output) > 0, "Reduce should produce non-empty output"


@pytest.mark.anyio
async def test_batch_node():
    """Test Batch node groups items correctly."""
    pipeline = load_template_bundle(Path("tests/pipelines/test_batch.soak"))
    pipeline.config.document_paths = [str(TEST_DATA)]
    pipeline.config.llm_credentials = LLMCredentials()

    result, error = await pipeline.run()
    assert error is None, f"Pipeline failed: {error}"

    chunks_node = [n for n in pipeline.nodes if n.name == "chunks"][0]
    batch_node = [n for n in pipeline.nodes if n.name == "batched"][0]

    # Batch should reduce number of items by grouping them
    assert batch_node.output is not None

    # Output is a BatchList with batches attribute
    assert hasattr(
        batch_node.output, "batches"
    ), "Batch output should have batches attribute"
    assert len(batch_node.output.batches) < len(
        chunks_node.output
    ), "Batch should group items and reduce count"

    # Each batch should be a tuple/list of items
    assert len(batch_node.output.batches) > 0, "Should have at least one batch"
    # Batches are tuples or lists of Box/dict objects from get_items()
    first_batch = batch_node.output.batches[0]
    assert hasattr(first_batch, "__iter__"), "Each batch should be iterable"
    batch_items = list(first_batch)
    assert len(batch_items) > 0, "Batch should contain items"
    assert isinstance(
        batch_items[0], (dict, object)
    ), "Batch items should be dicts or objects"


@pytest.mark.anyio
async def test_verify_quotes_node():
    """Test VerifyQuotes node validates quotes against source."""
    pipeline = load_template_bundle(Path("tests/pipelines/test_verify_quotes.soak"))
    pipeline.config.document_paths = [str(TEST_DATA)]
    pipeline.config.llm_credentials = LLMCredentials()

    result, error = await pipeline.run()
    assert error is None, f"Pipeline failed: {error}"

    verify_node = [n for n in pipeline.nodes if n.name == "verify"][0]

    # Should have quote matches
    assert verify_node.sentence_matches is not None
    assert len(verify_node.sentence_matches) > 0

    # Check statistics
    assert verify_node.stats is not None
    assert "n_quotes" in verify_node.stats
    assert "mean_bm25_score" in verify_node.stats
    assert "mean_cosine" in verify_node.stats

    # Each match should have required fields
    for match in verify_node.sentence_matches:
        assert "quote" in match
        assert "span_text" in match
        assert "bm25_score" in match
        assert "cosine_similarity" in match


@pytest.mark.anyio
async def test_pipeline_serialization():
    """Test that pipelines can be serialized to JSON."""
    pipeline = load_template_bundle(Path("tests/pipelines/test_map.soak"))
    pipeline.config.document_paths = [str(TEST_DATA)]
    pipeline.config.llm_credentials = LLMCredentials()

    result, error = await pipeline.run()
    assert error is None, f"Pipeline failed: {error}"

    # Test serialization
    json_output = pipeline.model_dump_json()
    assert len(json_output) > 0, "Pipeline should serialize to JSON"
    assert "nodes" in json_output, "Serialized pipeline should contain nodes"


if __name__ == "__main__":
    import asyncio

    async def run_all_tests():
        """Run all tests when executed directly."""
        print("Running Split node test...")
        await test_split_node()
        print("✓ Split node test passed\n")

        print("Running Map node test...")
        await test_map_node()
        print("✓ Map node test passed\n")

        print("Running Classifier node test...")
        await test_classifier_node_with_agreement()
        print("✓ Classifier node test passed\n")

        print("Running Filter node test...")
        await test_filter_node()
        print("✓ Filter node test passed\n")

        print("Running Transform node test...")
        await test_transform_node_multi_slot()
        print("✓ Transform node test passed\n")

        print("Running Reduce node test...")
        await test_reduce_node()
        print("✓ Reduce node test passed\n")

        print("Running Batch node test...")
        await test_batch_node()
        print("✓ Batch node test passed\n")

        print("Running VerifyQuotes node test...")
        await test_verify_quotes_node()
        print("✓ VerifyQuotes node test passed\n")

        print("Running serialization test...")
        await test_pipeline_serialization()
        print("✓ Serialization test passed\n")

        print("All tests passed!")

    asyncio.run(run_all_tests())
